在 Cloud TPU 上训练 PyTorch 模型

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

Cloud TPU 虚拟机

使用 Cloud TPU 虚拟机时,无需用户虚拟机,您的模型和代码将在 TPU 主机上的虚拟机中运行。您可以直接通过 SSH 连接到 TPU 主机。您可以运行任意代码安装软件包、查看日志、调试代码等。

每次运行新的 PyTorch/XLA 程序时,PyTorch/XLA 都会创建 TensorFlow 本地服务器。XRT (XLA 运行时) 客户端连接到本地 TensorFlow 服务器。

图片

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

旧版 TPU 节点

使用旧版 TPU 节点时,您需要为每个 TPU 添加用户虚拟机。用户虚拟机通过 gRPC 调用与 TPU 节点通信。TPU 节点运行 TensorFlow 服务器。XRT (XLA 运行时) 客户端通过网络连接到 TensorFlow 服务器。

图片