以下代码比较了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