13

It is currently possible to specify which CPU or GPU to use with the tf.device(...) function for specific ops, but is there anyway where you can specify a core of a CPU?

PhABC
  • 1,583
  • 1
  • 13
  • 19

1 Answers1

24

There's no API for pinning ops to a particular core at present, though this would make a good feature request. You could approximate this functionality by creating multiple CPU devices, each with a single-threaded threadpool, but this isn't guaranteed to maintain the locality of a core-pinning solution:

with tf.device("/cpu:4"):
  # ...

with tf.device("/cpu:7"):
  # ...

with tf.device("/cpu:0"):
  # ...

config = tf.ConfigProto(device_count={"CPU": 8},
                        inter_op_parallelism_threads=1,
                        intra_op_parallelism_threads=1)
sess = tf.Session(config=config)
mrry
  • 125,488
  • 26
  • 399
  • 400
  • 1
    Wow, just with a simple test the speed gain was 3 folds. Thank you so much! EDIT: After investigation, the 3 fold gain in speed is only caused by setting the config as described in your answer. It makes sens giving that I am using RNN. I will look into this! – PhABC Jun 16 '16 at 19:14
  • BTW, I think you need to set inter_op_parallelism_threads to higher numer in the example above, otherwise it will run everything sequentially, here's a test https://gist.github.com/yaroslavvb/b73ff35424dd7ab762234620cf583aac – Yaroslav Bulatov Sep 16 '16 at 23:09
  • 1
    I want only one thread. However, config = tf.ConfigProto(device_count={"CPU": 1}, inter_op_parallelism_threads=1, intra_op_parallelism_threads=1) generates 8 threads (One thread per core and I have 8 cores). How can I reduce the number of threads to 1? – fisakhan Aug 18 '20 at 08:38