配置开发环境¶
本指南的目标是在已安装 PyTorch/XLA 的 Cloud TPU 上设置一个交互式开发环境。如果您是首次使用 TPU，我们建议您从 Colab 和 Kaggle 开始。这两个选项都预装了 PyTorch/XLA 及其依赖项和生态系统包。有关最新的示例列表，请参阅我们的主要 README。
如果您想设置更定制化的开发环境，请继续阅读。
Visual Studio Code¶
前提条件
在您的本地计算机上安装了 Visual Studio Code 及 Remote Development 扩展
一个拥有 Cloud TPU 配额的 GCP 项目。有关申请 Cloud TPU 配额的更多信息，请参阅官方文档
在
ssh-agent中注册的 SSH 密钥。如果您尚未执行此操作，请参阅GitHub 的文档
在开始之前，请导出包含您拥有 Cloud TPU 配额的 GCP 项目和区域的环境变量
export PROJECT=...
export ZONE=...
export TPU_TYPE=... # e.g. "v2-8"
创建并连接到您的 TPU¶
创建一个已注册您的 SSH 密钥的 Cloud TPU VM
# Assuming your SSH key is named `id_ed25519`
gcloud compute tpus tpu-vm create --project=$PROJECT --zone=$ZONE --accelerator-type=$TPU_TYPE --version=tpu-ubuntu2204-base --metadata="ssh-keys=$USER:$(cat ~/.ssh/id_ed25519.pub)" $USER-tpu
检查您的 TPU 是否具有外部 IP 并通过 SSH 连接到它
gcloud compute tpus tpu-vm describe --project=$PROJECT --zone=$ZONE $USER-tpu --format="value(networkEndpoints.accessConfig.externalIp)"
# Output: 123.123.123.123
给您的 TPU 起一个友好的名字，以便后续步骤更轻松
echo -e Host $USER-tpu "\n " HostName $(gcloud compute tpus tpu-vm describe --project=$PROJECT --zone=$ZONE $USER-tpu --format="value(networkEndpoints.accessConfig.externalIp)") >> ~/.ssh/config
通过 SSH 连接到您的 TPU 以测试您的连接
ssh $USER-tpu
使用 PyTorch/XLA 设置 Visual Studio Code 工作区¶
从 VS Code 命令面板中，选择 `Remote-SSH: Connect to Host <https://vscode.js.cn/docs/remote/ssh>[__ 并选择您刚刚创建的主机（命名为 ]{.title-ref}[$USER-tpu]{.title-ref}`)。然后 VS Code 将打开一个连接到您的 TPU VM 的新窗口。
从内置的 Terminal 中，创建一个新文件夹作为工作区（例如 mkdir ptxla）。然后从 UI 或命令面板打开该文件夹。
注意：此时可选（但建议）安装官方的 Python 扩展，并通过命令面板（Python: Create Environment）创建一个 venv 虚拟环境。
安装最新的 PyTorch 和 PyTorch/XLA 版本
pip install numpy torch torch_xla[tpu] \
-f https://storage.googleapis.com/libtpu-wheels/index.html \
-f https://storage.googleapis.com/libtpu-releases/index.html
创建文件 test.py
import torch_xla as xla
# Optional
xla.runtime.set_device_type("TPU")
print("XLA devices:", xla.real_devices())
从您的终端运行测试脚本
$ python test.py
# Output: XLA devices: ['TPU:0', 'TPU:1', 'TPU:2', 'TPU:3', 'TPU:4', 'TPU:5', 'TPU:6', 'TPU:7']
# Number of devices will vary based on TPU type
