python - Tensorflow Eager 模式 : First execution on GPU slow

Question

下面的代码比较了 CPU 和 GPU 上的计算时间。只有第一次执行时，我在 GPU 上的运行时间比 CPU 慢，而在所有后续运行中，GPU 的速度要快得多。为什么第一次在 GPU 上运行很慢？如何让 GPU 上的第一次运行速度更快？

from __future__ import absolute_import, division, print_function

import tensorflow as tf

tf.enable_eager_execution()

import time

def time_matmul(x):
  start = time.time()
  for loop in range(10):
    tf.matmul(x, x)

  result = time.time()-start

  print("10 loops: {:0.2f}ms".format(1000*result))

print("On GPU:")
# Force execution on GPU #0 if available
if tf.test.is_gpu_available():
  with tf.device("GPU:0"): # Or GPU:1 for the 2nd GPU, GPU:2 for the 3rd etc.
    x = tf.random_uniform([1000, 1000])
    assert x.device.endswith("GPU:0")
    time_matmul(x)

# Force execution on CPU
print("On CPU:")
with tf.device("CPU:0"):
  x = tf.random_uniform([1000, 1000])
  assert x.device.endswith("CPU:0")
  time_matmul(x)

第一次运行的输出:

On GPU:
10 loops: 443.04ms
On CPU:
10 loops: 100.01ms

后续运行的输出:

On GPU:
10 loops: 1.00ms
On CPU:
10 loops: 103.01ms

PS:这和貌似related question不一样因为 tf.device("GPU:0") 已经选择了 /device:GPU:0 而不是 /device:XLA_GPU:0

Answer 1

出于好奇，我在 3 年后尝试了 OP 脚本。同样的情况发生在最新版本的 TF、CUDA(还是一张旧的 GTX1050 卡)上。一个可能的解释是数据移动。

在第一次运行时——无论是 GPU 还是 CPU——数据四处移动，为行动做好准备。众所周知，数据移动会显着降低速度。 CPU 内存在物理上比 GPU 内存“更近”，后者通常位于外部板上。默认计算是 CPU 及其内存，因此一个程序几乎可以为 CPU 运行做好准备——几乎不需要移动，并且基本上停留在同一个芯片上。 GPU 内存在物理上是一个不同的芯片，“距离很远”，因此移动到那里可能需要更多时间。

可以通过遍历 OP 脚本来支持这种想法(为了匹配 TF2.9.1 而稍作更改):

import tensorflow as tf

tf.compat.v1.enable_eager_execution()

import time

def time_matmul(run, x):
  start = time.time()
  for loop in range(10):
    tf.matmul(x, x)
  result = time.time()-start
  print(f"Run #{run}: {1000*result:0.2f}ms")

print("On GPU:")
# Force execution on GPU #0 if available
if tf.test.is_gpu_available():
  with tf.device("GPU:0"): # Or GPU:1 for the 2nd GPU, GPU:2 for the 3rd etc.
    x = tf.random.uniform([1000, 1000])
    assert x.device.endswith("GPU:0")
    for run in range(10):
        time_matmul(run, x)

# Force execution on CPU
print("On CPU:")
with tf.device("CPU:0"):
  x = tf.random.uniform([1000, 1000])
  assert x.device.endswith("CPU:0")
  for run in range(10):
      time_matmul(run, x)

结果是:

Run #0: 273.66ms
Run #1: 0.37ms
Run #2: 0.36ms
Run #3: 0.36ms
Run #4: 0.37ms
Run #5: 0.36ms
Run #6: 0.35ms
Run #7: 0.41ms
Run #8: 0.37ms
Run #9: 0.35ms
On CPU:
Run #0: 56.89ms
Run #1: 44.31ms
Run #2: 47.60ms
Run #3: 46.97ms
Run #4: 46.40ms
Run #5: 44.84ms
Run #6: 43.88ms
Run #7: 45.28ms
Run #8: 43.46ms
Run #9: 43.57ms

目测会发生什么(正确的统计方法会运行那么多次，完成但没有更多的洞察力)第一次运行很慢，但后来更快，更重要的是稳定。稳定性是我们首先期望的(运行相同应该表现相同)，但第一次运行需要通过将数据放在内存中“正确”的位置来设置。

我不知道有 API 可以手动放置数据，然后开始运行。但这将是一种“错觉”。这里的 Run #0 包括移动和计算。将两者分开可能会使运行 #0 与所有其他运行一样快，但我们仍然必须事先移动数据——所需时间不会显示在结果表中...

---

请注意，这种内存移动是一个可能的原因(此处归因推理)，并且可能还有其他事情发生。脚本结果支持这种想法，但它只是允许得出内存移动是可能原因的结论。这个帖子证明不了什么。正确分析以获得根本原因需要更多时间使用探查器(而 Python 探查器可能还不够)。

撇开这个免责声明不谈，它看起来确实像是我们在这里观察到的内存移动成本。