Create a per session inter op thread pool instead of relying on a global thread

pool: this ensures that the thread pool is reset whenever the session is reset.
Change: 153718067
This commit is contained in:
Benoit Steiner 2017-04-20 08:27:45 -08:00 committed by TensorFlower Gardener
parent 79d6be3b86
commit f12c967893

View File

@ -41,7 +41,10 @@ SingleMachine::SingleMachine(int timeout_s, int num_cpu_cores, int num_gpus)
}
CHECK_GE(num_cpu_cores, 1);
options_.config.set_intra_op_parallelism_threads(num_cpu_cores);
options_.config.set_inter_op_parallelism_threads(num_cpu_cores);
// Create a session specific thread pool to ensure the threads are reset when
// the session is reset.
options_.config.add_session_inter_op_thread_pool()->set_num_threads(
num_cpu_cores);
if (timeout_s > 0) {
options_.config.set_operation_timeout_in_ms(timeout_s * 1000);
}