TPU 比 GPU 慢？答案

【问题标题】：TPU slower than GPU?TPU 比 GPU 慢？
【发布时间】：2019-03-05 22:26:14
【问题描述】：

我刚刚尝试在 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()

【问题讨论】：

我怀疑这还不足以作为充分利用 TPU 的基准。您可以尝试调整示例笔记本中的微基准：colab.research.google.com/notebooks/…
这就是我适应上面示例代码的方式......虽然我首先发现了 MNIST 示例的问题（GPU 大约需要 0.8s/100 步，但 TPU 需要 2.2s/100 步）。我怀疑需要编写一些特定于 TPU 的代码才能利用它...github.com/tensorflow/tensorflow/tree/master/tensorflow/…
是的，您必须抽取 100 兆 MB 并循环多次才能看到差异。但它比 same 代码扩展更快更重要。你只需要设置从 1 个 TPU 到 N 个 TPU。在 8 之后使用 GPU，您必须开始进行 horovd 和集群分发并更改设置和...

标签： tensorflow machine-learning gpu google-colaboratory google-cloud-tpu

【解决方案1】：

正确地对设备进行基准测试很难，因此请对您从这些示例中学到的所有知识持保留态度。一般来说，比较您感兴趣的特定模型（例如运行 ImageNet 网络）以了解性能差异会更好。也就是说，我知道这样做很有趣，所以...

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

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

1 0.010800600051879883
10 0.09931182861328125
100 0.5581905841827393
500 2.7688047885894775

。所以你实际上可以在 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()

【讨论】：