配置开发环境

本指南的目的是在安装了 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

# Optional
torch_xla.runtime.set_device_type("TPU")

print("XLA devices:", torch_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

后续步骤

就是这样！您现在应该已经设置了一个安装了 PyTorch/XLA 的远程 Visual Studio Code 工作区。要运行更具代表性的示例，请参阅我们的 示例指南。