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

时间:2017-07-12 16:52:23

标签: c++ tensorflow

我正在编写一个围绕tensorflow 1.2 C API的c ++包装器(用于推理目的,如果重要的话)。 由于我的应用程序是多进程和多线程的,在显式分配资源的情况下,我想限制Tensorflow只使用一个线程。

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

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”消息。虽然我没有找到在此图分配之前指定选项的方法......

2 个答案:

答案 0 :(得分:3)

我遇到了同样的问题并通过在创建我的应用程序创建的第一个TF会话时设置线程数来解决它。如果没有使用选项对象创建第一个创建的会话,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来限制inter&amp;的数量。内部线程为1。

修改 重要说明:此代码在从主应用程序(谷歌样本培训师)调用时有效但当我将其移动到专用于包裹tensorFlow的DLL时停止工作。 TF 1.4.1忽略我传递的参数并旋转所有线程。 我想听听你的意见......

答案 1 :(得分:0)

使用TensorFlow C API没问题。 C API的限制是生成至少N个线程,其中N是内核数。您无法进一步减少它。

从环境中设置OMP_NUM_THREADS可以更改线程数,但是TensorFlow会覆盖这些设置并生成N个线程。

但是,您可以指定一个或多个要处理的核心。 tasket,numatcl和设置线程亲和性可以将给定进程锁定到给定核心,但不会更改线程数。

上述解决方案将减少TensorFlow生成的线程的总数但不减少到1 )。 TensorFlow生成的线程的总数仍将是多个,具体取决于CPU中的内核数。在大多数情况下,只有一个线程处于活动状态,而其他线程处于睡眠模式。我认为不可能有单线程 TensorFlow。

以下github问题支持了我的观点。 https://github.com/tensorflow/tensorflow/issues/33627 https://github.com/usnistgov/frvt/issues/30

从源代码构建TensorFlow C ++ API并更改源代码可能会有所帮助。