設定開發環境¶
本指南的目標是在 Cloud TPU 上設定安裝 PyTorch/XLA 的互動式開發環境。如果您是第一次使用 TPU，我們建議您從 Colab 和 Kaggle 開始。這兩個選項都預先安裝了 PyTorch/XLA 以及相依性和生態系統套件。如需最新的範例列表，請參閱我們主要的 README。
如果您想設定更客製化的開發環境，請繼續閱讀。
Visual Studio Code¶
先決條件
在本機上安裝具有 Remote Development 擴充功能 的 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.dev.org.tw/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
