配置开发环境¶
本指南的目标是在安装了 PyTorch/XLA 的 Cloud TPU 上设置交互式开发环境。如果这是您第一次使用 TPU，我们建议您从 Colab 和 Kaggle 开始。这两个选项都预装了 PyTorch/XLA 以及依赖项和生态系统软件包。有关最新示例列表，请参阅我们的主要 README。
如果您想设置更自定义的开发环境，请继续阅读。
Visual Studio Code¶
先决条件
Visual Studio Code 以及安装在本地计算机上的 远程开发扩展
具有 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。
从内置的 终端 中，创建一个新文件夹以用作工作区（例如 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
