Cloud Tensor Processing Unit(云张量处理单元)

张量处理单元 (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 可在以下地区使用:

  • 美国 (US)
    • us-central1-b
    • us-central1-c
    • us-central1-f(仅限 TFRC 计划
  • 欧洲 (EU)
    • europe-west4-a
  • 亚太地区 (APAC)
    • asia-east1-c

Cloud TPU 编程模型

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

以下是通过 TPUEstimator 实现的编程模型的一些显著特征:

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

与传统训练的区别

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

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

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

模型开发最佳做法

Cloud TPU 硬件包含一个矩阵单元 (MXU),旨在加速由密集矩阵乘法和卷积主导的程序。如果程序的执行时间中没有相当大的一部分用于执行矩阵乘法,可能不太适合 Cloud TPU。例如,一个程序的计算主要是非矩阵操作(例如添加、调整形状或连接),则可能无法实现高 MXU 利用率。以下指导原则可帮助您选择和构建适用于 Cloud TPU 的模型。

副本

一个 Cloud TPU 设备由四个芯片组成,每个芯片有两个 TPU 核心。因此,为了有效利用 Cloud TPU,程序应该使用全部八个核心。 TPUEstimator 使用图运算符来构建和运行复制计算。每个副本从本质上讲都是在每个核心上运行的训练图的一个复制品,并且训练包含整个批量大小 1/8 的小批量。

布局

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 操作列表。

何时使用 TPU

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

  • CPU:

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

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

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

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

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

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

后续步骤

  • 阅读 TPU 快速入门,开始使用 Cloud TPU 资源。
  • 完成 Cloud TPU 教程之一,了解如何在 Cloud TPU 资源上运行常见的机器学习工作负载。
此页内容是否有用?请给出您的反馈和评价:

发送以下问题的反馈:

此网页
Cloud TPU