python - 测量在 Tensorflow 中将数据从 RAM 移动到 GPU 内存所需的时间

Question

我想进行以下简单实验。

我正在使用 TensorFlow。我有一个大数组(5000x5000 float32 元素)。如何测量将此数组从 RAM 移动到 GPU 内存实际需要多长时间？

我知道我可以创建一些非常简单的计算图，运行它并测量它花费了多长时间。但是这有两个问题。首先，我担心测量的时间将主要取决于计算所花费的时间，而不是将数据从 RAM 转移到 GPU。其次，如果计算不涉及我提到的大数组，Tensorflow 将简化计算图，这样大数组就不会在其中，也根本不会从 RAM 移动到 GPU。

Answer 1

解决方案是做一个简单的基准测试，其中内存传输占主导地位。要检查 TensorFlow 是否没有优化您的传输，您可以在结果上添加一个小操作。像填充这样的微小操作的开销应该是几微秒，与将 100MB 加载到 GPU 相比，这是微不足道的，后者是 >5 毫秒。

def feed_gpu_tensor():
  params0 = create_array()
  with tf.device('/gpu:0'):
    params = tf.placeholder(tf.float32)
    result = tf.concat([params, tf.fill([1],1.0)], axis=0)
  for i in range(args.num_iters):
    with timeit('feed_gpu_tensor'):
      sess.run(result.op, feed_dict = {params: params0})

要运行这个基准测试，你可以这样做

wget https://github.com/diux-dev/cluster/blob/master/yuxin_numpy/tf_numpy_benchmark.py
python tf_numpy_benchmark.py --benchmark=feed_gpu_tensor

我发现在 p3.16xlarge 上，使用 tcmalloc(通过 LD_PRELOAD)，此副本 (100MB) 将花费 8 毫秒。

此外，作为健全性检查，您可以查看时间表。时间线将有 MEMCPYH2D 操作，它是实际的 CPU->GPU 副本，您可以使用它来确认它支配您的微基准测试步骤运行时

相关问题:

• D2H 和 H2D 基准测试:https://github.com/tensorflow/tensorflow/issues/17204

• 64 字节对齐输入数据:https://github.com/tensorflow/tensorflow/issues/17233