代码之家  ›  专栏  ›  技术社区  ›  Elliott Slaughter

在TensorFlow中使用多个CPU核心

  •  0
  • Elliott Slaughter  · 技术社区  · 7 年前

    我已经广泛研究了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核心,我必须对它做什么?

    注:

    • this answer 除其他外,我正在设置 设备计数 内部并行线程 .
    • 跟踪命令来自 this answer .
    • 我可以移除 tf.device 调用,它似乎对我的CPU利用率没有任何影响。

    我正在使用Conda安装的TensorFlow 1.10.0。

    1 回复  |  直到 7 年前
        1
  •  0
  •   Elliott Slaughter    7 年前

    在一段时间之后 TensorFlow issue here 我们确定问题在于程序被一个不断折叠的过程“优化”,因为输入都是琐碎的。结果是这个恒定的折叠过程按顺序运行。因此,如果您希望观察并行执行,那么实现这一点的方法是使输入变得非常重要,这样常量折叠就不会应用于它们。问题中建议的方法是使用 tf.placeholder ,我已经编写了一个示例程序,可以在这里使用它:

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

    程序输出样本见原版: https://github.com/tensorflow/tensorflow/issues/22619