如何减少 Tensorflow/Keras 使用的 CPU 数量？

Posted 2023-02-23

【中文标题】如何减少 Tensorflow/Keras 使用的 CPU 数量？

【英文标题】：How can I reduce the number of CPUs used by Tensorlfow/Keras?

【发布时间】：2020-01-15 09:56:59

【问题描述】：

我使用的是 Tensorflow 2.0 的 Keras api。

在我的 Keras 模型上调用 fit 时，它使用所有可用的 CPU。

我想限制使用的 CPU 数量。但是它在以前版本的 TensorFlow 中的工作方式不能再使用了：

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: 模块 'tensorflow.python.keras.api._v2.keras.backend' 没有属性 'set_session'

我该怎么做？

【参考方案1】：

在 TensorFlow 2.0 中，不再有会话。在 Eager Execution 中，直接使用 config API 在程序开始时设置并行度，如下所示。

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

【参考方案2】：

this post 中的第二个答案可能是限制使用的 CPU 数量的解决方案。您可以将代码更改为

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)

【讨论】：

TF 2.0 中没有 Session 和 ConfigProto