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

使用 gcloud 创建 Cloud TPU

在默认用户项目、网络和可用区中创建 TPU 虚拟机 运行：

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-8 \
   --version=tpu-ubuntu2204-base

创建 TPU 时，您可以将额外的 --network 和 --subnetwork 标记，如果 您要指定默认网络和子网。 如果您不想使用默认网络，则必须将 --network 标志。--subnetwork 标志是可选的，可用于 为您使用的任何网络指定默认子网（默认或 ）。请参阅gcloud 如需详细了解这些标志，请参阅 API 参考文档页面。

连接到 Cloud TPU 虚拟机

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central2-b

在 TPU 虚拟机上安装 PyTorch/XLA

   (vm)$ pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
   

设置 TPU 运行时配置

确保 PyTorch/XLA 运行时使用 TPU。

   (vm) $ export PJRT_DEVICE=TPU

执行简单的计算：

1. 在当前目录中创建一个名为 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)
   

2. 运行脚本：

     (vm)$ 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')