我刚刚尝试在Google Colab中使用TPU，我想看看TPU比GPU快多少。我惊讶地得到了相反的结果。
以下是NN。
random_image = tf.random_normal((100, 100, 100, 3))
result = tf.layers.conv2d(random_image, 32, 7)
result = tf.reduce_sum(result)
性能结果:
CPU: 8s
GPU: 0.18s
TPU: 0.50s
我不知道为什么。...TPU的完整代码如下:
def calc():
random_image = tf.random_normal((100, 100, 100, 3))
result = tf.layers.conv2d(random_image, 32, 7)
result = tf.reduce_sum(result)
return result
tpu_ops = tf.contrib.tpu.batch_parallel(calc, [], num_shards=8)
session = tf.Session(tpu_address)
try:
print('Initializing global variables...')
session.run(tf.global_variables_initializer())
print('Warming up...')
session.run(tf.contrib.tpu.initialize_system())
print('Profiling')
start = time.time()
session.run(tpu_ops)
end = time.time()
elapsed = end - start
print(elapsed)
finally:
session.run(tf.contrib.tpu.shutdown_system())
session.close()
最佳答案
正确地进行基准测试很难，因此，请您从这些示例中学到的一切都花点儿力气。通常最好比较您感兴趣的特定模型(例如，运行ImageNet网络)以了解性能差异。也就是说，我知道这样做很有趣，所以...
较大的模型将更好地说明TPU和GPU性能。您的示例还包括将编译时间包含在TPU调用的成本中:给定程序和形状的第一次调用之后的每个调用都将被缓存，因此除非您想捕获编译时间，否则您将需要在启动计时器之前先进行一次tpu_ops 。
当前，每次调用TPU函数之前，权重都将权重复制到TPU，然后才能开始运行，这会严重影响较小的操作。这是一个示例，该示例在返回CPU之前在TPU上运行循环，并显示以下输出。
。因此，您实际上可以在0.55秒内运行此函数的100次迭代。
import os
import time
import tensorflow as tf
def calc(n):
img = tf.random_normal((128, 100, 100, 3))
def body(_):
result = tf.layers.conv2d(img, 32, 7)
result = tf.reduce_sum(result)
return result
return tf.contrib.tpu.repeat(n[0], body, [0.0])
session = tf.Session('grpc://' + os.environ['COLAB_TPU_ADDR'])
try:
print('Initializing TPU...')
session.run(tf.contrib.tpu.initialize_system())
for i in [1, 10, 100, 500]:
tpu_ops = tf.contrib.tpu.batch_parallel(calc, [[i] * 8], num_shards=8)
print('Warming up...')
session.run(tf.global_variables_initializer())
session.run(tpu_ops)
print('Profiling')
start = time.time()
session.run(tpu_ops)
end = time.time()
elapsed = end - start
print(i, elapsed)
finally:
session.run(tf.contrib.tpu.shutdown_system())
session.close()