有没有办法在GPU上使用tensorflow map_fn？

Question

我有一个形状[a,n] 的张量A,我需要my_op用另一个形状[b,n]的张量B执行一个op ,使得得到的张量C具有形状[a,b].

换句话说:对于A中的每个子指标(A [0],A 1,... A [n]),我需要对B中的每个子指标执行元素操作.

因此产生的张量将包含以下内容:

[ [ A[0] op B[0] , A[0] op B[1], ... , A[0] op B[b] ],
  [ A[1] op B[0] , A[1] op B[1], ... , A[1] op B[b] ],
  [ ...                                             ],
  [ A[a] op B[0] , A[a] op B[1], ... , A[a] op B[b] ] ]

我能够找到的唯一方法就是通过嵌套使用tf.map_fn 这样:

import tensorflow as tf
import time
import numpy as np

a_size = 64
b_size = 256*256
n = 256
A = tf.placeholder(tf.float32,[a_size,n])
B = tf.placeholder(tf.float32,[b_size,n])

def elementwise_op(a,b):
    return tf.reduce_sum(tf.multiply(a,b))

def intermediate_op(sub_a,my_b):
    sample_values = tf.map_fn(lambda x: elementwise_op(sub_a,x),my_b)
    return sample_values

my_op = tf.map_fn(lambda x: intermediate_op(x,B),A)

with tf.Session() as sess:
    a = np.random.rand(a_size,n)
    b = np.random.rand(b_size,n)
    start_time = time.time()
    result = sess.run (my_op,feed_dict={A:a,B:b})
    print ("exec time: " ,time.time()-start_time)
    print (result.shape)

上面的代码运行正常,但是它没有很好地使用GPU(根据使用率只有~15%nvidia-smi).事实上,当仅使用CPU时,它的运行速度提高了一个数量级!(在我的12核机器上)当使用GPU运行时,我看到我的一个 CPU内核的GPU利用率非常低(~15%)和100%.仅在CPU上运行时,我看到所有CPU内核的利用率为100%.

仅运行5个CPU的平均时间:11.33秒

5 GPU运行的平均时间:111.88s

上述测试使用官方Tensorflow docker镜像运行:( tensorflow/tensorflow:latest-py3对于CPU)和tensorflow/tensorflow:latest-gpu-py3(对于GPU)

我的猜测是map_fn,通过python lambda,强制数据在每次迭代时在CPU和GPU之间来回复制,而op的嵌套特性使它变得更糟.在未回答的太问题的意见在这里表明,这是如此.

本文声称:

lambda表达式是GPU利用率低的主要原因.

-

所以我的问题是:有没有办法强制map_fn使用GPU？还是要避免Python lambda？

或者,是否有其他(可能更多tensorflow-y)方法来实现上述结果,以便获取图形在GPU上运行？

编辑: 运行探查器之后(我必须大幅减少数组的大小才能让探测器完全运行,因为它像疯了一样占用RAM),以下几行引起了我的注意:

node name     |     output bytes     |      total execution time     | accelerator execution time     |     cpu execution time

Mul                    1.02KB (22.23%, 0.29%),      195.07ms (85.00%, 13.06%),       5.29ms (100.00%, 25.79%),      189.78ms (84.79%, 12.89%)

Sum                      256B (21.41%, 0.07%),      241.48ms (69.08%, 16.17%),        6.01ms (74.21%, 29.29%),      235.47ms (69.01%, 15.99%)

TensorArrayScatterV3      512B (0.64%, 0.15%),      658.31ms (46.87%, 44.09%),        9.19ms (44.80%, 44.80%),      649.12ms (46.90%, 44.08%)

看起来某些操作主要在CPU上进行,并且只在一个线程上完成!

Answer 1

该tf.map_fn()构造可以与在 GPU 上运行操作的函数一起使用。默认情况下，TensorFlow 将尝试在 GPU 上运行尽可能多的函数，任何与 GPU 不兼容的操作都将在 CPU 上运行。在您的程序中，整个elementwise_op()函数是由 GPU 兼容的操作构建的，因此每次迭代时 CPU 和 GPU 之间不应该有额外的复制。

GPU利用率低的原因很难从程序片段中确定。例如，如果A和B相对较小，并且您从 Python 提供它们并立即取回结果，则将初始数据复制到 GPU 或从 GPU 复制初始数据的开销可能会占主导地位。追踪此问题的最佳方法是使用 GPU 分析器，您可以使用tfprof该分析器或NVIDIA Visual Profiler。