2

下面的代码比较了 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：这与看似相关的问题不同，因为tf.device("GPU:0")已经选择/device:GPU:0而不是/device:XLA_GPU:0

4

0 回答 0