在TensorFlow中使用多个CPU内核

时间:2018-09-25 23:07:30

标签: python multithreading tensorflow parallel-processing affinity

我已经在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设备:< / p>

/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_countinter_op_parallelism_threads听起来都像我想要的,但似乎都没有真正导致使用多个内核。我尝试了所有我能想到的组合,包括将一个或另一个设置为1以防它们相互冲突,并且似乎没有任何作用。

我还可以用taskset确认我对CPU的亲和力没有做任何奇怪的事情:

$ taskset -p $$
pid 21395's current affinity mask: ffffffffff

要使该代码使用多个CPU内核,我到底需要做什么?

注意:

  • this answer中,我正在设置device_countinter_op_parallelism_threads
  • 跟踪命令来自this answer
  • 我可以删除tf.device调用,这似乎对我的CPU使用率没有任何影响。

我正在使用从conda安装的TensorFlow 1.10.0。

1 个答案:

答案 0 :(得分:1)

TensorFlow issue here上来回往复后,我们确定问题在于程序正在通过不断的折叠传递来“优化”,因为输入都是微不足道的。事实证明,这种恒定的折页传递顺序进行。因此,如果要观察并行执行,执行此操作的方法是使输入变得平凡,以便常量折叠不会应用于它们。问题中建议的方法是使用tf.placeholder,并且我在这里编写了一个示例程序来使用它:

https://gist.github.com/elliottslaughter/750a27c832782f4daec8686281027de8

有关该程序的示例输出,请参见原始问题:https://github.com/tensorflow/tensorflow/issues/22619