0

I have one tensorflow inference task in C++, e.g. class TaskA, with inter_op_parallelism_threads=1 and intra_op_parallelism_threads=1 to use single thread.

When TaskA is actually used, I will create 2 threads {t1, t2} to use two models. I want all tensorflow calculation pinned to certain CPU, such as t1 on CPU:0 and t2 on CPU:3.

I try to pin certain thread to TaskA below:

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);
    }
}

But this can just pin UseTaskA() function to certain CPU, tensorflow will create its own thread and may use other CPU.

Is there any way to pin tensorflow's session run to certain cpu?

Notice: my problem is not the same as "how to run tensorflow on CPU". It's actually how to bind each tensorflow session to certain CPU when create multi sessions.

user10878853
  • 9
  • 2

0 Answers0