PyTorch を使用した Cloud TPU での Resnet50 のトレーニング

このチュートリアルでは、PyTorch を使用して Cloud TPU デバイスで ResNet-50 モデルをトレーニングする方法を説明します。PyTorch と ImageNet データセットを使用する、TPU 用に最適化されたその他のイメージ分類モデルにも、同じパターンを適用できます。

このチュートリアルのモデルは、残余ネットワーク（ResNet）アーキテクチャを最初に導入する画像認識のためのディープ残余ラーニングに基づいています。このチュートリアルでは、50 層のバリアントの ResNet-50 を使用して、PyTorch/XLA を使ったモデルのトレーニング方法を説明します。

目標

  • データセットを準備します。
  • トレーニング ジョブを実行します。
  • 出力結果を確認します。

費用

このドキュメントでは、課金対象である次の Google Cloudコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新規の Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、 Google Cloud プロジェクトが正しく設定されていることを確認します。

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

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

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

    6. このチュートリアルでは、 Google Cloudの課金対象コンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金が発生しないように、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

TPU VM を作成する

  1. Cloud Shell ウィンドウを開きます。

    Cloud Shell を開く

  2. TPU VM を作成する

    gcloud compute tpus tpu-vm create your-tpu-name \
--accelerator-type=v3-8 \
--version=tpu-ubuntu2204-base \
--zone=us-central1-a \
--project=your-project

  3. SSH を使用して TPU VM に接続します。

    gcloud compute tpus tpu-vm ssh  your-tpu-name --zone=us-central1-a

  4. TPU VM に PyTorch/XLA をインストールします。

    (vm)$ pip install torch torch_xla[tpu] torchvision -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html

  5. PyTorch/XLA GitHub リポジトリのクローンを作成する

    (vm)$ git clone --depth=1 https://github.com/pytorch/xla.git

  6. 架空のデータでトレーニング スクリプトを実行する

    (vm) $ PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

  1. TPU VM との接続を解除します。

    (vm) $ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. TPU VM を削除します。

    $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=us-central1-a

次のステップ