我已经广泛研究了TensorFlow上的其他答案,但我似乎无法让它在我的CPU上使用多个内核。
根据HTOP,以下程序只使用一个CPU内核:
import tensorflow as tf
n_cpus = 20
sess = tf.Session(config=tf.ConfigProto(
device_count={ "CPU": n_cpus },
inter_op_parallelism_threads=n_cpus,
intra_op_parallelism_threads=1,
))
size = 100000
A = tf.ones([size, size], name="A")
B = tf.ones([size, size], name="B")
C = tf.ones([size, size], name="C")
with tf.device("/cpu:0"):
x = tf.matmul(A, B)
with tf.device("/cpu:1"):
y = tf.matmul(A, C)
sess.run([x, y])
# run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
# run_metadata = tf.RunMetadata()
# sess.run([x, y], options=run_options, run_metadata=run_metadata)
# for device in run_metadata.step_stats.dev_stats:
# device_name = device.device
# print(device.device)
# for node in device.node_stats:
# print(" ", node.node_name)
但是,当我取消对底部的行的注释并更改
size
因此,计算实际上在合理的时间内完成,我看到TensorFlow似乎认为它至少使用了2个CPU设备:
/job:localhost/replica:0/task:0/device:CPU:0
_SOURCE
MatMul
_retval_MatMul_0_0
_retval_MatMul_1_0_1
/job:localhost/replica:0/task:0/device:CPU:1
_SOURCE
MatMul_1
从根本上讲,我想在这里并行地对不同的核心执行不同的操作。我不想在多个核心上拆分一个操作,尽管我知道在这个人为的例子中这是可行的。两个
device_count
和
inter_op_parallelism_threads
听起来像我想要的,但这两个都不可能真正导致使用多个内核。我已经尝试了所有我能想到的组合,包括设置一个或另一个
1
以防它们互相冲突,似乎什么都不起作用。
我也可以确认
taskset
我对我的CPU亲和力没有做任何奇怪的事情:
$ taskset -p $$
pid 21395's current affinity mask: ffffffffff
为了让这段代码使用多个CPU核心,我必须对它做什么?
注:
我正在使用Conda安装的TensorFlow 1.10.0。