使用 PyTorch 在 Cloud TPU 虚拟机上运行计算

本文档简要介绍了如何使用 PyTorch 和 Cloud TPU。

准备工作

在运行本文档中的命令之前，您必须创建 Google Cloud 账号、安装 Google Cloud CLI 并配置 gcloud 命令。如需了解详情，请参阅设置 Cloud TPU 环境。

使用 `gcloud` 创建 Cloud TPU

定义一些环境变量，以便更轻松地使用命令。

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east5-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

环境变量说明

变量	说明
`PROJECT_ID`	您的 Google Cloud 项目 ID。使用现有项目或创建新项目。
`TPU_NAME`	TPU 的名称。
`ZONE`	要在其中创建 TPU 虚拟机的可用区。如需详细了解支持的可用区，请参阅 TPU 区域和可用区。
`ACCELERATOR_TYPE`	加速器类型用于指定您要创建的 Cloud TPU 的版本和大小。如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。
`RUNTIME_VERSION`	Cloud TPU 软件版本。

通过运行以下命令创建 TPU 虚拟机：

$ gcloud compute tpus tpu-vm create $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --accelerator-type=$ACCELERATOR_TYPE \
    --version=$RUNTIME_VERSION

连接到 Cloud TPU 虚拟机

使用以下命令通过 SSH 连接到 TPU 虚拟机：

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

如果您无法使用 SSH 连接到 TPU 虚拟机，可能是因为 TPU 虚拟机没有外部 IP 地址。如需访问没有外部 IP 地址的 TPU 虚拟机，请按照连接到没有公共 IP 地址的 TPU 虚拟机中的说明操作。

在 TPU 虚拟机上安装 PyTorch/XLA

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

验证 PyTorch 是否可以访问 TPU

使用以下命令验证 PyTorch 是否可以访问 TPU：

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

该命令的输出应如下所示：

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

执行基本计算

在当前目录中创建一个名为 tpu-test.py 的文件，并将以下脚本复制粘贴到其中：

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

运行脚本：

(vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

脚本的输出显示计算结果：

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

清理

为避免因本页中使用的资源导致您的 Google Cloud 账号产生费用，请按照以下步骤操作。

断开与 Cloud TPU 实例的连接（如果您尚未这样做）：
```
(vm)$ exit
```
您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

删除您的 Cloud TPU。

$ gcloud compute tpus tpu-vm delete $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

通过运行以下命令来验证资源已删除。确保您的 TPU 不再列出。删除操作可能需要几分钟时间才能完成。
```
$ gcloud compute tpus tpu-vm list \
    --zone=$ZONE
```

后续步骤

详细了解 Cloud TPU 虚拟机：