在 Cloud TPU 上运行 TensorFlow

本文档详细介绍了如何在 Cloud TPU 上运行 TensorFlow 模型。在 Cloud TPU 上训练 TensorFlow 模型的方式取决于您使用的 Cloud TPU 架构。如需了解详情,请参阅系统架构

当您运行应用时,TensorFlow 会生成计算图并将其发送到 TPU。TPU 主机及时编译计算图并将程序二进制文件发送到一个或多个 TPU 设备以供执行。模型的输入通常存储在 Cloud Storage 中。TPU 主机会将输入数据流式传输到一个或多个 TPU 设备以供训练。

Cloud TPU 虚拟机

使用 Cloud TPU 虚拟机架构时,没有用户虚拟机,您的模型和代码在 TPU 主机上的虚拟机中运行。您可以通过 SSH 直接连接到 TPU 主机。您可以在同一虚拟机或不同虚拟机上创建 TensorFlow 客户端和服务。您可以在 TPU 主机上运行任意代码以执行预处理等操作。

映像

如需详细了解如何在 TPU 虚拟机上训练 TensorFlow 模型,请参阅 Cloud TPU 虚拟机 Tensorflow 快速入门

Cloud TPU 节点

在 Cloud TPU 节点架构中,TPU 主机由 Compute Engine 虚拟机和 TPU 节点组成。Compute Engine 虚拟机运行 TensorFlow 客户端和模型。虚拟机通过 gRPC 与 TPU 节点通信。TPU 节点运行 TensorFlow 服务器和 XLA 编译器。

以下框图展示 Cloud TPU 节点软件架构,它由神经网络模型、TensorFlow 客户端、TensorFlow 服务器和 XLA 编译器组成。

图片

Cloud TPU TensorFlow API

Keras 是一个使用 Python 编写的高级深度学习 API,它在机器学习平台 TensorFlow 上运行。如果要编写使用 Cloud TPU 的神经网络模型,我们建议使用 Keras。

TPU Estimator(旧版)

TPU Estimator 是一套基于 Estimator 构建的旧版高级 API,可简化 Cloud TPU 的构建模型,并使 TPU 发挥最佳性能。虽然 Tensorflow 1.x 仍支持 TPU Estimator,但我们建议改用 Keras。

TensorFlow 服务器

TensorFlow 服务器在 Cloud TPU 服务器上运行。当 TensorFlow 服务器从 TensorFlow 客户端接收计算图时,它会执行以下操作:

  1. 从 Cloud Storage 加载输入
  2. 将图划分为可以在 Cloud TPU 上运行的部分以及必须在 CPU 上运行的部分
  3. 生成对应于要在 Cloud TPU 上运行的子图的 XLA 操作
  4. 调用 XLA 编译器

如需详细了解如何在 TPU 节点上训练 TensorFlow 模型,请参阅 Cloud TPU 快速入门