How can I reduce the number of CPUs used by Tensorlfow/Keras?

Question

I am using the Keras api of Tensorflow 2.0.

When calling fit on my Keras model, it uses all availabel CPUs.

I would like to limit the number of used CPUs. However the way it used to work in former versions of Tensorflow cannot be used anymore:

tf.keras.backend.set_session(tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(
       intra_op_parallelism_threads=2, inter_op_parallelism_threads=2)))

AttributeError: module 'tensorflow.python.keras.api._v2.keras.backend' has no attribute 'set_session'

How could I do that?

Answer 1

In Tensorflow 2.0, there is no session anymore. In eager execution, directly use the config API to set the parallelism at the start of the program like this.

import tensorflow as tf

tf.config.threading.set_intra_op_parallelism_threads(2)
tf.config.threading.set_inter_op_parallelism_threads(2)
with tf.device('/CPU:0'):
    model = tf.keras.models.Sequential([...

https://www.tensorflow.org/api_docs/python/tf/config/threading