大多数人认为，今天的深度学习仍基于英伟达的 GPU，几乎没有可替代的相关实际选择。谷歌专门为深度学习定制的芯片张量处理单元（Tensor Processing Unit，TPU）有望改变这种情况。
首次发布九个月后，谷歌终于在谷歌云平台上向早期测试版用户发布了 TPUv2。我们进行了几个快速基准测试。下面，我们分享一下这次评测的经验和初步结果。
我们一直希望深度学习硬件市场的竞争越来越激烈，日益激烈的竞争有望打破英伟达对深度学习硬件的垄断。除此之外，这还将定义未来深度学习的基础架构。
但要记住：正如谷歌在许多地方提到的，TPU 仍处于早期测试阶段，所以我们讨论的一些事情未来可能发生变化。
谷歌云上的 TPU
第一代芯片 TPUv1 面向推断（inference），而第二代芯片的重点在于加速学习。TPUv2 的内核中有一个广泛应用于深度学习的脉动阵列，负责执行矩阵乘法。Jeff Dean 在报告中称，每个 Cloud TPU 设备具备四个「TPUv2 芯片」。每个芯片有 16GB 内存和两个内核，每个内核中有两个矩阵乘法单元。两个内核一共提供 45 万亿次浮点运算（45 TFLOPs），整个 TPU 设备可提供 180 TFLOPs、64GB 内存。而目前这一代 Nvidia V100 GPU 可提供 125 TFLOPs、16GB 内存。
要在谷歌云平台上使用 TPU，你需要启动 Cloud TPU（在获得 quota 命令后执行此操作）。无需（或无法）将 Cloud TPU 分配给特定的 VM 实例，通过网络即可发现实例中的 Cloud TPU。每个 Cloud TPU 都被分配了一个名称和 IP 地址，你需要将二者提供给 TensorFlow 代码。
创建一个新的 Cloud TPU。请注意，Cloud TPU 有 IP 地址。
TPU 仅接受 TensorFlow 1.6 版本（发行候选版）的支持。除此之外，由于与 TPU 通信所需的各种代码都由 TensorFlow 提供，所以你的 VM 实例上不需要任何驱动程序。在 TPU 上执行的代码经过优化，且通过 TensorFlow 中的 XLA 进行实时编译。
为了高效使用 TPU，你的代码应该建立在高级别的 Estimator 抽象上。然后，你可以使用 TPUEstimator 来执行必要任务，提高使用 TPU 的效率，如为 TPU 设置数据排队并在不同核上进行并行计算。当然有使用 TPUEstimator 的方法，但目前我们还不知道例子或文档。
完成所有设置后，即可正常地运行 TensorFlow 代码，你会在启动过程中发现 TPU，计算图被编译并传输至 TPU。有趣的是，TPU 可以直接从云存储中读取和写入，以存储检验点（checkpoint）或事件摘要（event summary）。为了实现这一点，你需要提供支持 Cloud TPU 写入云存储的服务账户。
基准测试
有趣的地方不在于 TPU 的速度有多快。TensorFlow 有一个 TPU 模型的 GitHub 库，效果很好。下面，我们介绍在 TPU 上运行 ResNet 和 Inception 的实验。我们想看一个没有针对 TPU 进行优化的模型是如何运作的，因此我们在 TPU 上运行一个基于 LSTM 的文本分类模型。通常，谷歌推荐使用较大的模型，而这个模型较小，所以看看这种情况下 TPU 是否仍然具有优势很有意思。
对于所有的模型，我们将单个 Cloud TPU 的训练速度与单个 Nvidia P100 和 V100 GPU 进行了对比。我们注意到，一个完全的对比除了吞吐量之外，还应包括模型的最终质量和收敛性能。我们的实验算是最初的尝试，将对未来的研究留下详细的分析。
TPU 和 P100 的实验在谷歌云平台上的 n1-standard-16 实例（16 vCPU Intel Haswell，60 GB 内存）上运行。对于 V100 GPU，我们在 AWS 上使用 p3.2xlarge 实例（8 vCPU，60 GB 内存），所有系统都运行 Ubuntu 16.04。对于 TPU，我们从 PyPi 存储库安装了 TensorFlow 1.6.0-rc1。GPU 实验则使用 nvidia-docker 运行，nvidia-docker 使用包含 CUDA 9.0 和 cuDNN 7.0 的 TensorFlow 1.5 镜像（tensorflow: 1.5.0-gpu-py3）。
TPU 优化模型
我们先来看看正式针对 TPU 进行优化的模型的性能。下图中，你可以看到它们每秒处理图像的性能。
TPU 的批量大小是 1024，GPU 是 128。对于 GPU，我们使用了 TensorFlow 基准库的实现，使用 flag 『use_fp16=true』而不是标记『fp32』运行。左边两组条状对比的是混合精度的训练。训练数据是谷歌存储在云存储（用于 TPU）和本地磁盘（用于 GPU）上的伪 Imagenet 数据集。
在 ResNet-50 上，单个 Cloud TPU（包含 4 个 TPUv2 和 64GB 内存）的速度约是单个 P100 的 7.3 倍，V100 的 2.8 倍。对于 InceptionV3，加速几乎相同（分别为〜7.6 和〜2.5）。V100 具备更高精度（fp32），因此速度稍慢。注意：在 TPU 上训练精度为 fp32 的模型完全不可能，因为它只支持混合精度计算。
显然，除了速度之外，还必须考虑价格。该表显示按秒计费的按需定价的标准化性能，TPU 仍然明显领先。
自定义 LSTM 模型
我们的自定义模型是一个具备 1024 个隐藏单元的双向 LSTM 文本分类模型。LSTM 是目前 NLP 的一个基本构造块，所以这与基于计算机视觉的官方模型形成了很好的对比。
源代码已经使用了 Estimator 框架，所以调整它来使用 TPUEstimator 非常简单。但是一定要注意：在 TPU 上，我们无法让模型收敛，而在 GPU 上的相同模型（批量大小等相同）运行得很好。我们认为这是由于我们的代码或 TensorFlow 中存在 bug。由于模型无法收敛，这里我们决定不报告初步结果（因此删去了上一版中的初步结果），而是介绍我们在另一篇博文中的发现。
结论
在我们测试的模型上，TPU 与最新一代的 GPU 相比性能更好，也更加便宜。这与之前的报道（https://www.forbes.com/sites/moorinsights/2018/02/13/google-announces-expensive-cloud-tpu-availability/#d1ec931359f1）形成了鲜明对比。总之，就测试版而言。使用 TPU 和适应 TensorFlow 代码的经验已经很不错了。
我们认为，一旦 TPU 可以面向更多用户，它们就可以真正替代英伟达的 GPU。
原文链接：https://blog.riseml.com/benchmarking-googles-new-tpuv2-121c03b71384
