云张量处理单元 (TPU)

张量处理单元 (TPU) 是 Google 定制开发的专用集成电路 (ASIC)，用于加速机器学习工作负载。Google 结合机器学习领域的丰富经验与领先优势，从零开始打造了这些 TPU。

借助 Cloud TPU，您将能够使用 TensorFlow 在 Google 的 TPU 加速器硬件上运行机器学习工作负载。Cloud TPU 旨在实现最高性能与灵活性，帮助研究人员、开发者和企业构建可使用 CPU、GPU 和 TPU 的 TensorFlow 计算集群。高级 TensorFlow API 可帮助您在 Cloud TPU 硬件上运行模型。

TPU 的优点

Cloud TPU 资源提高了机器学习应用中大量使用的线性代数计算的性能。在训练大型复杂的神经网络模型时，TPU 可以最大限度地缩短达到准确率所需的时间。以前在其他硬件平台上需要花费数周时间进行训练的模型，在 TPU 中只需数小时即可收敛。

可用情况

Cloud TPU 地区和区域显示可以使用 Cloud TPU 的位置。

Cloud TPU 编程模型

在执行密集矢量和矩阵计算方面，Cloud TPU 的速度非常快。 与计算速度相比，在 Cloud TPU 和主机内存之间传输数据的速度较慢 - PCIe 总线的速度比 Cloud TPU 互连和片上高带宽存储器 (HBM) 慢得多。模型的部分编译（在这种情况下，执行作业在主机和设备之间之间折返）导致 TPU 在大多数时候处于空闲状态，等待数据通过 PCIe 总线到达。为了缓解这种情况，Cloud TPU 的编程模型专为在 TPU 上执行大部分训练而设计，理想情况可执行整个训练循环。

以下是 TPU 编程模型的一些显著特征：

  • 所有模型参数都保存在片上高带宽存储器中。
  • 在 Cloud TPU 上启动计算的成本通过在一个循环中执行多个训练步骤来摊销。
  • 输入训练数据被流式传输到 Cloud TPU 上的“馈入”队列。在 Cloud TPU 上运行的程序会在每个训练步骤中从这些队列中检索批次。
  • 在宿主机（连接到 Cloud TPU 设备的 CPU）上运行的 TensorFlow 服务器提取数据并对其进行预处理，然后再将其“馈入”Cloud TPU 硬件。
  • 数据同时载入：Cloud TPU 上的核心以同步方式执行驻留在各自 HBM 中的相同程序。所有核心在每个神经网络步结束时执行归约操作。

何时使用 TPU

Cloud TPU 针对特定工作负载进行了优化。在某些情况下，您可能希望在 Compute Engine 实例上使用 GPU 或 CPU 来运行机器学习工作负载。通常，您可以根据以下准则确定哪种硬件最适合您的工作负载：

  • CPU

    • 需要最高灵活性的快速原型设计
    • 训练时间不长的简单模型
    • 有效批量大小较小的小型模型
    • 以 C++ 编写的自定义 TensorFlow 操作占主导的模型
    • 受主机系统可用 I/O 或网络带宽限制的模型

  • GPU

    • 源代码不存在或更改起来太费力的模型
    • 具有大量自定义且必须至少部分在 CPU 上运行的 TensorFlow 操作的模型
    • 具有不能在 Cloud TPU 上使用的 TensorFlow 操作的模型（请参阅可用的 TensorFlow 操作列表）
    • 有效批量大小较大的中到大型模型

  • TPU

    • 由矩阵计算主导的模型
    • 在主训练循环内没有自定义 TensorFlow 操作的模型
    • 需要训练数周或数月的模型
    • 有效批量大小非常大的大型和极大型模型

Cloud TPU 不适合以下工作负载：

  • 需要频繁分支或逐项 (element-wise) 代数主导的线性代数程序。TPU 经过专门优化，可执行快速、庞大的矩阵乘法，因此不是由矩阵乘法主导的工作负载不太可能在 TPU 上有较好的表现（与其他平台相比）。
  • 以稀疏方式访问内存的工作负载可能不适用于 TPU。
  • 需要高精度算法的工作负载。例如，双精度算法不适用于 TPU。
  • 包含用 C++ 编写的自定义 TensorFlow 操作的神经网络工作负载。具体而言，主训练循环体中的自定义操作不适用于 TPU。

神经网络工作负载必须能够在 TPU 上运行整个训练循环的多次迭代。虽然这不是 TPU 本身的基本要求，但这是提高效率所需的 TPU 软件生态系统的当前限制之一。

与传统训练的区别

典型的 TensorFlow 训练图由多个提供各种功能的相互重叠的子图组成，包括：

  • 读取训练数据的 I/O 操作。
  • 输入预处理阶段，通常通过队列连接。
  • 模型变量。
  • 这些变量的初始化代码。
  • 模型本身。
  • 损失函数。
  • 梯度代码（通常自动生成）。
  • 用于监控训练的摘要操作。
  • 检查点的保存/恢复操作。

在 Cloud TPU 上，TensorFlow 程序由 XLA 即时编译器编译。在 Cloud TPU 上进行训练时，能够在硬件上编译和执行的代码只有对应于模型的密集部分、损失和梯度子图的代码。TensorFlow 程序的所有其他部分作为常规分布式 TensorFlow 会话在主机（Cloud TPU 服务器）上运行。这通常包括读取训练数据的 I/O 操作、所有预处理代码（例如解码压缩图像、随机采样/剪裁、组合训练小批量）以及图表的所有维护部分（如检查点保存/恢复）。

模型开发最佳做法

单个 Cloud TPU 芯片包含 2 个核心，每个核心包含多个矩阵单元 (MXU)，这是为了加速由密集矩阵乘法和卷积主导的程序（请参阅系统架构）。耗费大量时间执行矩阵乘法的程序通常非常适合 Cloud TPU。由非矩阵运算（如 add、reshape 或 concatenate）主导计算的程序可能无法实现较高的 MXU 利用率。以下指导原则可帮助您选择和构建适用于 Cloud TPU 的模型。

一个 Cloud TPU 设备由四个芯片组成，每个芯片有两个 TPU 核心。因此，为了有效利用 Cloud TPU，程序应该使用全部八个核心。

布局

XLA 编译器执行代码转换，包括将矩阵乘法平铺成较小的块，以便高效地在矩阵单元 (MXU) 上执行计算。XLA 编译器使用 MXU 硬件结构（128x128 脉动阵列）和 TPU 内存子系统设计（该设计更擅长处理 8 的倍数的维度）来提高平铺效率。因此，某些布局更有利于平铺，而另一些布局则需要在平铺之前调整形状。调整形状的操作在 Cloud TPU 上通常受限于内存。

形状

XLA 编译器即时为第一批量编译 TensorFlow 图。 如果任何后续批量具有不同形状，则该模型不起作用。 （在每次形状变化时重新编译图太慢。） 因此，如果模型张量形状在运行时动态变化，则该模型不适合 TPU。

填充

高性能的 Cloud TPU 程序可以将密集计算轻松地平铺为 128x128 的块。如果矩阵计算不能占满整个 MXU，编译器会用零来填充张量。填充有两个不足：

  • 用零填充的张量不能充分利用 TPU 核心。
  • 填充增加了张量所需的片上内存存储，并且在极端情况下可能导致内存不足错误。

虽然填充是由 XLA 编译器在必要时自动执行的，但您可以通过 op_profile 工具确定填充量。您可以通过选择非常适合 TPU 的张量维度来避免填充。

维度

选择合适的张量维度对于使 TPU 硬件（尤其是 MXU）发挥最高性能很有帮助。XLA 编译器尝试使用批量大小或特征维度来最大限度地利用 MXU。因此，这两个因素之一必须是 128 的倍数。否则，编译器会将其中某个填充到 128。理想情况下，批量大小和特征维度应为 8 的倍数，这样可以使内存子系统发挥最高性能。

运维

请参阅可用的 TensorFlow 操作列表。

VPC Service Controls 集成

借助 Cloud TPU VPC Service Controls，您可以为 Cloud TPU 资源定义安全边界，并控制跨边界的数据移动。如需详细了解 VPC Service Controls，请参阅 VPC Service Controls 概览。如需详细了解将 Cloud TPU 与 VPC Service Controls 结合使用的限制，请参阅支持的产品和限制。

Edge TPU

在云端训练的机器学习模型越来越需要“在边缘”（即在物联网 (IoT) 边缘运行的设备上）进行推断。这些设备包括传感器和其他智能设备，它们收集实时数据，作出智能的决策，然后采取行动或将信息传递给其他设备或云端。

由于此类设备必须在有限的电量（包括电池供电）下运行，因此 Google 设计了 Edge TPU 协处理器来加速低功耗设备上的机器学习推断。一个 Edge TPU 每秒可执行 4 万亿次操作（4 TOPS），能耗仅 2 瓦特，换句话说，每瓦特可获得 2 TOPS。例如，Edge TPU 能够以低能耗的方式以接近每秒 400 帧的速率执行先进的移动视觉模型，例如 MobileNet V2。

这款低功耗的机器学习加速器增强了 Cloud TPU 和 Cloud IoT 的功能，提供一个端到端（云端到边缘、硬件 + 软件）的基础架构，帮助您实现基于 AI 的解决方案。

Edge TPU 适用于多种设备类型的的原型设计和生产设备，包括单板计算机、模块化系统、PCIe/M.2 卡和表面安装模块。如需详细了解 Edge TPU 和所有可用产品，请访问 coral.ai。

后续步骤

  • 阅读 TPU 快速入门，开始使用 Cloud TPU 资源。
  • 完成 Cloud TPU 教程之一，了解如何在 Cloud TPU 资源上运行常见的机器学习工作负载。