c++ - 使用 C API 更改 Tensorflow 推理的线程数

Question

我正在围绕 tensorflow 1.2 C API 编写一个 c++ 包装器(用于推理目的，如果它重要的话)。由于我的应用程序是一个多进程和多线程的应用程序，资源是明确分配的，所以我想将 Tensorflow 限制为仅使用一个线程。

目前，运行一个允许批处理的简单推理测试，我看到它正在使用所有 CPU 内核。我已经尝试使用 C 和 C++ 的混合来限制新 session 的线程数，如下所示(请原谅我的部分代码片段，我希望这是有道理的):

tensorflow::ConfigProto conf;
conf.set_intra_op_parallelism_threads(1);
conf.set_inter_op_parallelism_threads(1);
conf.add_session_inter_op_thread_pool()->set_num_threads(1);
std::string str;
conf.SerializeToString(&str);
TF_SetConfig(m_session_opts,(void *)str.c_str(),str.size(),m_status);
m_session = TF_NewSession(m_graph, m_session_opts, m_status);

但我看不出它有什么不同 - 所有核心仍然得到充分利用。

我是否正确使用了 C API？

(我目前的解决方法是重新编译 Tensorflow，并将硬编码线程数设置为 1，这可能会奏效，但显然这不是最佳方法...)

-- 更新--

我也尝试添加:

conf.set_use_per_session_threads(true);

没有成功。仍然使用多核...

我还尝试以高日志详细程度运行，并得到了这个输出(仅显示我认为相关的内容):

tensorflow/core/common_runtime/local_device.cc:40] Local device intraop parallelism threads: 8
tensorflow/core/common_runtime/session_factory.cc:75] SessionFactory type DIRECT_SESSION accepts target: 
tensorflow/core/common_runtime/direct_session.cc:95] Direct session inter op parallelism threads for pool 0: 1

当我使用 TF_NewGraph() 实例化一个新图形时，“parallelism threads: 8”消息就会出现。虽然我没有找到一种方法来在此图分配之前指定选项...

Answer 1

我遇到了同样的问题，并通过在创建我的应用程序正在创建的第一个 TF session 时设置线程数解决了这个问题。如果第一个创建的 session 不是使用选项对象创建的，TF 将创建工作线程作为机器上的核心数 * 2。

这是我使用的 C++ 代码:

// Call when application starts
void InitThreads(int coresToUse)
{
    // initialize the number of worker threads
    tensorflow::SessionOptions options;
    tensorflow::ConfigProto & config = options.config;
    if (coresToUse > 0)
    {
        config.set_inter_op_parallelism_threads(coresToUse);
        config.set_intra_op_parallelism_threads(coresToUse);
        config.set_use_per_session_threads(false);  
    }
    // now create a session to make the change
    std::unique_ptr<tensorflow::Session> 
        session(tensorflow::NewSession(options));
    session->Close();
}

传递 1 以将内部和内部线程的数量限制为每个 1。

编辑:重要说明:此代码在从主应用程序(谷歌示例培训师)调用时有效，但在我将其移至专用于包装 tensorFlow 的 DLL 时停止工作)。 TF 1.4.1 忽略我传递的参数并启动所有线程。我想听听您的意见...