TF2.8 CPU及GPU速度测试

测试CPU及GPU在TF下面的速度差异

import tensorflow as tf
import time

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

  result = time.time()-start

  print("{} loops: {:0.2f}ms".format(range_1,1000*result))
    
print("Normal Running:")
x = tf.random.uniform([1000, 1000])
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)

# Force execution on GPU #0 if available
if tf.config.list_physical_devices("GPU"):
  print("On GPU:")
  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)```

--------
测试结果：
1000 loops: 372.44ms
On CPU:
1000 loops: 5694.99ms
On GPU:
1000 loops: 305.03ms