在 Cloud TPU 上运行 TensorFlow

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

概览

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

以下图表展示 Cloud TPU 软件架构,包括神经网络模型、TPU Estimator 和 TensorFlow 客户端、TensorFlow 服务器和 XLA 编译器。

图片

Cloud TPU TensorFlow API

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

TPU Estimator(旧版)

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

TensorFlow 客户端

TPU Estimator 将您的程序转换为 TensorFlow 操作,然后这些操作由 TensorFlow 客户端转换为计算图。TensorFlow 客户端将计算图传送给 TensorFlow 服务器。

TensorFlow 服务器

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

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

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

Cloud TPU 虚拟机

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

图片

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