Keras / tensorflow - 限制核心数（intra_op_parallelism_threads 不工作）

Question

我一直在尝试在 CPU 集群上运行 keras，为此我需要限制使用的内核数量（这是一个共享系统）。所以为了限制核心数量，我登陆了this answer。然而，这根本行不通。我尝试使用以下基本代码运行：

from keras.applications.vgg16 import VGG16
from keras import backend as K
import numpy as np

conf = K.tf.ConfigProto(device_count={'CPU': 1}, 
                        intra_op_parallelism_threads=2, 
                        inter_op_parallelism_threads=2)
K.set_session(K.tf.Session(config=conf))
model = VGG16(weights='imagenet', include_top=False)
x = np.random.randn(1000, 224, 224, 3)
features = model.predict(x)

当我运行它并检查 htop 时，它使用了所有 (128) 个逻辑核心。这是keras中的错误吗？还是我做错了什么？

Keras 说我的 CPU 支持 SSE4.1 和 SSE4.2，因为我没有从二进制编译，所以没有使用它们。从二进制编译也会解决原来的问题吗？

编辑：我在从 unix 机器启动 keras 脚本时发现了一种解决方法：

taskset -c 0-23 python keras_script.py

这将在机器的前 24 个内核上运行脚本。它可以工作，但如果它可以从 keras/tensorflow 中获得，它仍然会很好。

Answer 1

我发现这段代码对我有用，希望对你有帮助：

from keras import backend as K
import tensorflow as tf
jobs = 2 # it means number of cores
config = tf.ConfigProto(intra_op_parallelism_threads=jobs,
                         inter_op_parallelism_threads=jobs,
                         allow_soft_placement=True,
                         device_count={'CPU': jobs})
session = tf.Session(config=config)
K.set_session(session)