时间戳

首页 > 解决方案 > 如何在 C++ 中将一个 tensorflow 线程固定到一个特定的 CPU？

c++ - 如何在 C++ 中将一个 tensorflow 线程固定到一个特定的 CPU？

问题描述

我在 C++ 中有一个张量流推理任务，例如 TaskA 类，其中 inter_op_parallelism_threads=1 和 intra_op_parallelism_threads=1 使用单线程。

实际使用TaskA时，我会创建2个线程{t1, t2}来使用两个模型。我希望将所有 tensorflow 计算固定到某个 CPU，例如 CPU:0 上的 t1 和 CPU:3 上的 t2。

我尝试将某些线程固定到下面的 TaskA：

void UseTaskA() {
    task = TaskA();
    task.Run();  // similiar to session.Run()
}
int main() {
    std::vector<int> cpu_index = {0, 3};
    int num_threads = 2;
    for (int i = 0; i < num_threads; ++i) {
        std::thread t(UseTaskA);
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(cpu_index[i], &cpuset);
        pthread_setaffinity_np(t.native_handle(), sizeof(cpu_set_t), &cpuset);
    }
}

但这只能将 UseTaskA() 函数固定到某个 CPU，tensorflow 将创建自己的线程并可能使用其他 CPU。

有没有办法将 tensorflow 的会话运行固定到某个 cpu？

注意：我的问题与“如何在 CPU 上运行 tensorflow”不同。它实际上是在创建多个会话时如何将每个 tensorflow 会话绑定到某个 CPU。

标签： c++multithreadingtensorflowcpu

解决方案

推荐阅读