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

美国

TPU 类型 (v2) TPU v2 核心数 TPU 总内存 可用地区
v2-8 8 64 GiB us-central1-b
us-central1-c
us-central1-f
v2-32 32 256 GiB us-central1-a
v2-128 128 1 TiB us-central1-a
v2-256 256 2 TiB us-central1-a
v2-512 512 4 TiB us-central1-a
TPU 类型 (v3) TPU v3 核心数 TPU 总内存 可用地区
v3-8 8 128 GiB us-central1-a
us-central1-b
us-central1-f

欧洲

TPU 类型 (v2) TPU v2 核心数 TPU 总内存 可用地区
v2-8 8 64 GiB europe-west4-a
v2-32 32 256 GiB europe-west4-a
v2-128 128 1 TiB europe-west4-a
v2-256 256 2 TiB europe-west4-a
v2-512 512 4 TiB europe-west4-a
TPU 类型 (v3) TPU v3 核心数 TPU 总内存 可用地区
v3-8 8 128 GiB europe-west4-a
v3-32 32 512 GiB europe-west4-a
v3-64 64 1 TiB europe-west4-a
v3-128 128 2 TiB europe-west4-a
v3-256 256 4 TiB europe-west4-a
v3-512 512 8 TiB europe-west4-a
v3-1024 1024 16 TiB europe-west4-a
v3-2048 2048 32 TiB europe-west4-a

亚太地区

TPU 类型 (v2) TPU v2 核心数 TPU 总内存 可用地区
v2-8 8 64 GiB 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 中的相同程序。所有核心在每个神经网络步结束时执行归约操作。

何时使用 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 软件生态系统的当前限制条件之一。

与传统训练的区别

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

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

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

模型开发最佳做法

单个云 TPU 芯片包含 2 个核心,每个核心包含多个矩阵单元 (MXU),这是为了加速由密集矩阵乘法和卷积主导的程序(请参阅硬件架构)。 耗费大量时间执行矩阵乘法的程序通常非常适合 Cloud TPU。由非矩阵运算(如 add、reshape 或 concatenate)主导计算的程序可能无法实现较高的 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 操作列表。

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 资源上运行常见的机器学习工作负载。