使用 C API 更改张量流推理的线程数

Change number of threads for Tensorflow inference with C API

本文关键字：推理 线程 张量流 API 使用      更新时间：2023-10-16

我正在围绕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(( 实例化新图形，就会显示"并行线程：8"消息。不过，我没有找到在此图形分配之前指定选项的方法......

我遇到了同样的问题，并通过在创建应用程序正在创建的第一个 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 将线程间和线程内的数量限制为每个 1 个。

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

使用 TensorFlow C API 没有问题。C API 的限制是生成至少 N 个线程，其中 N 是内核数。你不能进一步减少它。

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

但是，您可以指定一个或多个要处理的内核。 任务集、Numatcl 和 设置线程关联可以将给定进程锁定到给定核心，但不会更改线程数。

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

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

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

相关文章：