Cloud TPU 简介

张量处理单元 (TPU) 是 Google 专为机器学习工作负载设计的硬件加速器。如需详细了解 TPU 硬件,请参阅系统架构。Cloud TPU 是一种 Web 服务,使 TPU 可用作 Google Cloud Platform 上的可扩缩计算资源。

TPU 使用专为执行机器学习算法中常见的大型矩阵运算而设计的硬件,更高效地训练模型。TPU 具有片上高带宽内存,可让您使用较大的模型和批量大小。TPU 可以连接在名为 Pod 的组中,这让您可以在几乎不需要更改代码的情况下对工作负载进行纵向扩容。

我们建议您使用单个 TPU 启动机器学习项目,并扩容到 TPU Pod 以用于生产环境。首先,您可以利用 Google 研究和工程团队优化用于 TPU 的一组开源参考模型。如需了解详情,请参阅参考模型

工作原理

为了理解 TPU 的工作原理,我们不妨了解一下其他加速器如何应对训练机器学习模型的计算挑战。

CPU 的工作方式

CPU 是一种基于冯·诺依曼结构的通用处理器。这意味着 CPU 与软件和内存协同工作,如下所示:

CPU 工作方式图示

CPU 最大的优点是它们的灵活性。您可以在 CPU 上为许多不同类型的应用加载任何类型的软件。例如,您可以使用 CPU 在 PC 上进行文字处理、控制火箭发动机、处理银行交易或通过神经网络对图片进行分类。

对于每次计算,CPU 从内存加载值,对值执行计算,然后将结果存储回内存中。与计算速度相比,内存访问速度较慢,并可能会限制 CPU 的总吞吐量。这通常称为冯·诺依曼瓶颈

GPU 的工作方式

为了提高吞吐量,GPU 在单个处理器中包含数千个算术逻辑单元 (ALU)。现代 GPU 通常包含 2500 - 5000 个 ALU。如此大量的处理器意味着您可以同时执行数千次乘法和加法运算。

GPU 工作方式图示

这种 GPU 架构非常适合并行处理大量运算(例如神经网络中的矩阵运算)的应用。实际上,在用于深度学习的典型训练工作负载上,GPU 的吞吐量可比 CPU 高出一个数量级。

不过,GPU 仍然是一种通用处理器,必须支持许多不同应用和软件。也就是说,GPU 与 CPU 存在同样的问题。对于数千个 ALU 中的每一次计算,GPU 都必须访问寄存器或共享内存,以读取运算对象以及存储中间计算结果。

TPU 的工作方式

Google 设计了 Cloud TPU,它们是专门用于神经网络工作负载的矩阵处理器。TPU 不能运行文字处理程序、控制火箭引擎或执行银行交易,但它们可以很快地处理神经网络中使用的大量矩阵运算。

TPU 的主要任务是矩阵处理,这是乘法和累加运算的组合。TPU 包含数千个乘法累加器,这些累加器彼此直接连接以形成大型物理矩阵。这称为脉动阵列架构。对于 Cloud TPU v3,单个处理器上有两个 128 x 128 ALU 的脉动阵列。

TPU 主机将数据流式传输到馈入队列中。TPU 从馈入队列加载数据,并将其存储在 HBM 内存中。计算完成后,TPU 会将结果加载到馈出队列中。然后,TPU 主机从馈出队列读取结果并将其存储在主机的内存中。

为了执行矩阵操作,TPU 将 HBM 内存中的参数加载到 MXU 中。

图示:TPU 如何从内存加载参数

然后,TPU 从内存加载数据。每次执行乘法运算时,系统都会将结果传递给下一个乘法累加器。输出是数据和参数之间所有乘法结果的总和。在矩阵乘法过程中,不需要访问内存。

图示:TPU 如何从内存加载数据

因此,TPU 可以在神经网络计算中实现高计算吞吐量。

何时使用 TPU

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

CPU

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

GPU

  • 具有大量自定义且必须至少部分在 CPU 上运行的 TensorFlow/PyTorch/JAX 操作的模型
  • 具有不能在 Cloud TPU 上使用的 TensorFlow/PyTorch 操作的模型
  • 有效批量大小较大的中到大型模型

TPU

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

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

  • 需要频繁分支或包含大量元素级代数的线性代数程序
  • 以稀疏方式访问内存的工作负载
  • 需要高精度算法的工作负载
  • 主训练循环中包含自定义操作的神经网络工作负载

VPC Service Controls 集成

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

Cloud TPU 使用入门

设置 Google Cloud 帐号 您必须先创建 Google Cloud 帐号和项目,然后才能使用 Cloud TPU 资源。
激活 Cloud TPU API 要训练模型,您必须激活 Compute Engine 和 Cloud TPU API。
向 Cloud TPU 授予对您的 Cloud Storage 存储分区的访问权限 Cloud TPU 需要访问存储数据集的 Cloud Storage 存储桶。
训练模型 阅读 Cloud TPU 快速入门教程之一,以开始使用 Cloud TPU。
分析模型 在模型训练过程中,使用 TensorBoard 或其他工具直观呈现模型并跟踪关键指标,例如学习速率、损失和准确率。

后续步骤

想要详细了解 Cloud TPU?以下资源可能对您有所帮助。

快速入门:使用 Compute Engine 阅读我们的快速入门之一,尝试使用 Cloud TPU 训练模型。
TPU Colab 使用免费的 Colab 体验 Cloud TPU。
Cloud TPU 教程 使用参考机器学习模型测试 Cloud TPU。
Cloud TPU 系统架构 详细了解 TPU。
价格 了解 Cloud TPU 如何以经济高效的方式处理您的机器学习工作负载。
与销售人员联系 想要探讨特定的实现或应用?与 Cloud TPU 销售代表联系。