使用 PyTorch 在 Cloud TPU 上訓練 ResNet-50

注意：本頁面適用於 Cloud TPU API，但該 API 已停止積極開發。Cloud TPU API 只會收到錯誤修正和安全性更新。如要使用最新功能並支援最新 TPU 版本，建議您使用 Compute Engine 或 Google Kubernetes Engine 管理 TPU 資源。詳情請參閱「Compute Engine 中的 TPU 資源」和「GKE 中的 TPU 簡介」。TPU7x 和更新的 TPU 版本不支援 Cloud TPU API。

這個教學課程將說明如何使用 PyTorch 在 Cloud TPU 裝置中訓練 ResNet-50 模型。如果您有其他已針對 TPU 完成最佳化處理的圖片分類模型，而且這些模型使用的是 PyTorch 和 ImageNet 資料集，您也可以按照這個教學課程中的步驟對其進行訓練。

本教學課程中的模型是以圖像識別的深度殘差學習為基礎，該論文首度提出殘差網路 (ResNet) 架構的概念。這個教學課程使用了含有 50 層架構的變化版本「ResNet-50」，並說明如何使用 PyTorch/XLA 訓練模型。

警告：本教學課程使用的是第三方資料集。Google 不對這個資料集的有效性或其他方面提供任何聲明、擔保或其他保證。

目標

• 準備資料集。
• 執行訓練工作。
• 驗證輸出結果。

費用

在本文件中，您會使用下列 Google Cloud的計費元件：

• Compute Engine
• Cloud TPU

如要根據預測用量估算費用，請使用 Pricing Calculator。

初次使用 Google Cloud 的使用者可能符合免費試用期資格。

事前準備

開始學習這個教學課程之前，請先檢查 Google Cloud 專案設定是否正確。

登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎 建立帳戶，親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額，可用於執行、測試及部署工作負載。

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

本逐步操作說明使用 Google Cloud的計費元件，請參閱 Cloud TPU 定價頁面來估算費用。使用完畢後，請務必清除您建立的資源，以免產生不必要的費用。

建立 TPU VM

1. 開啟 Cloud Shell 或終端機視窗。

   開啟 Cloud Shell

2. 定義 Google Cloud CLI 指令參數的環境變數：

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-central1-b
   export ACCELERATOR_TYPE=v6e-8
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   環境變數說明

   • PROJECT_ID： 您的 Google Cloud 專案 ID。使用現有專案或建立新專案。
   • TPU_NAME：TPU 的名稱。
   • ZONE： 要建立 TPU VM 的可用區。如要進一步瞭解支援的區域，請參閱 TPU 地區和區域。
   • ACCELERATOR_TYPE： 加速器類型會指定要建立的 Cloud TPU 版本和大小。如要進一步瞭解各個 TPU 版本支援的加速器類型，請參閱 TPU 版本。
   • RUNTIME_VERSION：Cloud TPU 軟體版本。

3. 建立 TPU VM

    gcloud compute tpus tpu-vm create $TPU_NAME \
    --accelerator-type=$ACCELERATOR_TYPE \
    --version=$RUNTIME_VERSION \
    --zone=$ZONE \
    --project=$PROJECT_ID
   

   注意： 首次在新的 Cloud Shell VM 中執行指令時，系統會顯示 Authorize Cloud Shell 頁面。按一下頁面底部的 Authorize，允許 gcloud 使用您的憑證發出 API 呼叫。 Google Cloud

4. 使用 SSH 連線至 TPU VM：

    gcloud compute tpus tpu-vm ssh  $TPU_NAME --zone=$ZONE
   

在 TPU VM 上安裝 PyTorch/XLA

1. 建立名為 requirements.txt 的檔案，並加入下列內容：

   --find-links https://storage.googleapis.com/libtpu-releases/index.html
   --find-links https://storage.googleapis.com/libtpu-wheels/index.html
   torch~=2.6.0
   torch_xla[tpu]~=2.6.0
   torchvision
   ray[default]==2.40.0
   

2. 安裝依附元件：

   pip install -r requirements.txt
   

3. 複製 PyTorch/XLA GitHub 存放區

   git clone --depth=1 https://github.com/pytorch/xla.git
   

4. 使用偽造資料執行訓練指令碼

   PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py \
   --fake_data \
   --batch_size=256 \
   --num_epochs=1
   

清除所用資源

為避免因為本教學課程所用資源，導致系統向 Google Cloud 帳戶收取費用，請刪除含有相關資源的專案，或者保留專案但刪除個別資源。

1. 中斷與 TPU VM 的連線：

   exit
   

2. 刪除 TPU VM。

   gcloud compute tpus tpu-vm delete $TPU_NAME \
       --zone=$ZONE
   

後續步驟

• 使用 PyTorch 訓練擴散模型
• 排解在 TPU 上執行 PyTorch 時的問題
• PyTorch/XLA 說明文件