9

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?

kederrac
  • 16,819
  • 6
  • 32
  • 55
Pop
  • 12,135
  • 5
  • 55
  • 68

2 Answers2

9

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

Manoj Mohan
  • 5,654
  • 1
  • 17
  • 21
0

The second answer in this post could be a solution to limit the number of CPUs used. You may change the code as 

import tensorflow as tf
from keras import backend as K

num_cores = 4

num_CPU = 1
num_GPU = 0

config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
                        inter_op_parallelism_threads=num_cores, 
                        allow_soft_placement=True,
                        device_count = {'CPU' : num_CPU,
                                        'GPU' : num_GPU}
                       )

session = tf.Session(config=config)
K.set_session(session)
aesari
  • 207
  • 1
  • 6