配置开发环境¶
本指南旨在为安装了 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
