Tensorflow急切模式:在GPU上的第一次执行

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

出于好奇，三年后我尝试了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可以手动放置数据，然后开始运行。但这将是一种“幻觉”。这里的运行#0包括移动和计算。将两者分开可能会使运行#0与所有其他运行一样快，但是我们仍然必须提前移动数据--所需的时间将不会显示在结果表中.

请注意，这个记忆运动是一个可能的原因(这里的绑架推理)，并且可能有其他的事情发生。这种想法得到了脚本结果的支持，但它只允许得出记忆运动是一个可能的原因。这篇文章证明不了什么。正确的分析以获得根本原因将需要更多的时间使用分析器(而Python分析器可能还不够)。

撇开这个免责声明不谈，我们在这里观察到，这确实是一个内存移动的代价。