TPU比GPU慢？

Question

我刚刚尝试在Google中使用TPU，我想看看TPU比GPU快多少。我意外地得到了相反的结果。

以下是神经网络。

  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()

Answer 1

正确的基准测试设备是很难的，所以请把你从这些例子中学到的一切都拿出来。一般来说，最好比较您感兴趣的特定模型(例如运行ImageNet网络)，以了解性能差异。我知道这么做很有趣所以..。

较大的模型将更好地说明TPU和GPU的性能。您的示例还包括TPU调用成本中的编译时间:对于给定程序和形状的第一个调用之后的每个调用都将被缓存，因此除非您想要捕获编译时间，否则您希望在启动计时器之前缓存一次tpu_ops。

目前，每个对TPU函数的调用都会在TPU开始运行之前将权重复制到TPU，这对小型操作的影响更大。下面是一个示例，它在TPU上运行一个循环，然后返回到CPU，输出如下。

• 1 0.010800600051879883
• 10 0.09931182861328125
• 100 0.5581905841827393
• 2.7688047885894775

。因此，您实际上可以在0.55s中运行这个函数的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()