Use single thread on Intel CPU


(Sbuschjaeger) #1

Hi,

I have been playing around with the C++ Frontend for PyTorch on my Laptop (Intel® Core™ i7-4600U) and were able to include PyTorch into my CPP app by following the MNIST example (https://github.com/goldsborough/examples/blob/cpp/cpp/mnist/mnist.cpp).
My app already utilizes parallelization to some degree, so I would like to run training / inference on a single thread. However, I was not able to tell PyTorch that - it would always use all cores. Here is what I did so far:

  • I could not find any “torch.set_num_threads(1)” so I searched the PyTorch sources for anything related and found at::set_num_threads(1);
  • Setting the openmp environment variable “OMP_NUM_THREADS=1”
  • Setting the MKL environment variable “MKL_NUM_THREADS=1”
  • Further investigation of the source code revealed, that caffee2 uses a ThreadPool which is initialized with cpuinfo_get_processors_count() from the cpuinfo lib (https://github.com/pytorch/cpuinfo). I did not find a way to set this from the outside

note, that I did not compile PyTorch myself, but used the library provided on the website. There seem to be a lot of different frameworks involved such as openmp,mkl,mkldnn etc. Thus I am a little confused on how to force PyTorch to use one thread. Any ideas?

Thanks!