Tensorflow 的 while 循环在 GPU 上运行缓慢？

Question

由于未知原因，以下代码在 GPU 上比在 CPU 上慢了两倍。谁能解释为什么：

import time
import tensorflow as tf

with tf.device('/device:GPU:0'):  # gpu takes: 5.132448434829712 seconds
    # with tf.device('/cpu:0'): # cpu takes: 3.440524101257324 seconds
    i = tf.constant(0)
    while_condition = lambda i: tf.less(i, 2 ** 20)
    a = tf.fill([16, 16], 1.1)
    b = tf.fill([16, 16], 2.2)
    def body(i):
        res = tf.matmul(a, b)
        # increment i
        add = tf.add(i, 1)

        return (add,)


    ini_matmul = tf.matmul(a, b)

    # do the loop:
    loop = tf.while_loop(while_condition, body, [i])

with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    sess.run(ini_matmul)  # force GPU to initilise anything it needs.

    t0 = time.time()
    sess.run(loop)

    t1 = time.time()
    print(t1 - t0)
sess.close()

注意：通常GPU运行5秒，CPU运行3秒，使用numpy的CPU版本仅运行1.5秒。硬件：在 Google 的 Colab 上运行的 Tensorflow 代码。在本地英特尔酷睿 i5-7267U 上运行的 Numpy 代码。

Numpy 版本：

import numpy as np
import time

i = 0
a = np.full([16,16],1.1)
b = np.full([16,16],2.2)

t0 = time.time()

while i < 2**20:
    a.dot(b)
    i += 1

t1 = time.time()

print(t1-t0)

更新

这对我来说变得越来越紧密，因为扩大矩阵并没有真正的帮助。这是其中的更新代码和数据（运行 Titan XP 卡/Intel i7 CPU）。基本上tensorflow运行速度要慢得多。

import time
import tensorflow as tf

dimension = 11
repeat = 2**10
use_gpu = False
# Device: /device:GPU:0, Dimension 11, Repeat: 1024, Time cost: 0.00457597 seconds.
# Device: /cpu:0, Dimension 11, Repeat: 1024, Time cost: 0.00353599 seconds.

dev_name = '/device:GPU:0' if use_gpu else '/cpu:0'

with tf.device(dev_name):  
    i = tf.constant(0)
    while_condition = lambda i: tf.less(i, repeat)
    a = tf.constant(1.1, shape=[2**dimension, 2**dimension])
    b = tf.constant(2.2, shape=[2**dimension, 2**dimension])
    def body(i):
        res = tf.matmul(a, b)
        add = tf.add(i, 1)
        return (add,)
    ini_matmul = tf.matmul(a, b)
    # do the loop:
    loop = tf.while_loop(while_condition, body, [i])

with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    sess.run(ini_matmul)  # force initialisation.

    t0 = time.time()
    sess.run(loop)
    t1 = time.time()
    print('Device: {dev}, Dimension {dim:d}, Repeat: {r:d}, Time cost: {t:.8f} seconds.'.format(
        dev = dev_name,
        dim = dimension, r = repeat,
        t = t1 - t0
    ))
sess.close()

Answer 1

这是个有趣的问题。

您在 TensorFlow 代码片段中看到的 GPU 和 CPU 执行之间的相对减慢几乎肯定是由于 GPU内存分配开销造成的。总结一下链接，cudaMalloc比 慢malloc。当且仅当加速超过内存分配时间的差异时，这种内存分配减慢才会被请求操作的加速所抵消（matmul在本例中） 。当矩阵很大时总是如此。当矩阵很小时，情况就不是这样了，就像你的例子中的情况一样。为了验证这个假设，迭代地增加被乘数的大小并记录 CPU 和 GPU 的运行时间 - 如果内存分配确实是问题所在，则两者应该收敛，然后交叉。matmul

Numpy 运行时间和仅 CPU 运行时间之间的差异可能是由于 Numpy 和 TensorFlow 代码之间非常细微的差异造成的。请注意，在 Numpy 代码中，您仅实例化a一次b。看起来您在 TensorFlow 代码中做了同样的事情，因为您只调用了一次初始化，但您仍然在每次迭代中填充张量！要了解原因，请注意tf.fill返回一个Tensor. 根据定义，Tensor每次sess.run在包含对象的图形上调用对象时都会填充对象。因此，这两个片段实际上做的事情略有不同。更直接的比较是在 TensorFlow 代码片段中进行a和。btf.constant