在 TensorFlow 中使用多个 CPU 核心

我广泛研究了有关 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 and inter_op_parallelism_threads听起来像我想要的，但似乎都没有真正导致使用多核。我已经尝试了所有我能想到的组合，包括将其中一个或另一个设置为1万一它们相互冲突，并且似乎没有任何作用。

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

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

我到底需要对这段代码做什么才能让它使用多个 CPU 核心？

Note:

• From 这个答案其中我正在设置device_count and inter_op_parallelism_threads.
• 跟踪命令来自这个答案.
• 我可以删除tf.device调用，它似乎对我的 CPU 利用率没有任何影响。

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

经过一番来回TensorFlow 问题在这里我们确定问题在于该程序正在通过恒定的折叠传递进行“优化”，因为输入都很微不足道。事实证明，这个恒定的折叠过程是按顺序运行的。因此，如果您想观察并行执行，执行此操作的方法是使输入变得不平凡，以便常量折叠不会应用于它们。该问题中建议的方法是使用tf.placeholder，我在这里编写了一个利用它的示例程序：

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

请参阅原始问题以获取程序的示例输出：https://github.com/tensorflow/tensorflow/issues/22619