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

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

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

目標

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

費用

このチュートリアルでは、Google Cloud Platform の課金対象となる以下のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを出すことができます。初めて GCP を使用される場合は、無料トライアルをご利用いただけます。

始める前に

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

  1. Google アカウントにログインします。

    Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。

  2. GCP Console のプロジェクト セレクタのページで、GCP プロジェクトを選択または作成します。

    プロジェクト セレクタのページに移動

  3. Google Cloud Platform プロジェクトに対して課金が有効になっていることを確認します。 プロジェクトに対して課金が有効になっていることを確認する方法を学習する

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

VM を作成する

コンソール に移動して、新しい VM インスタンスを作成し、VM の名前とゾーンを指定します。

実際のデータで Resnet50 をトレーニングする場合は、できるだけ CPU 数の多いマシンタイプを選択してください。Resnet50 は通常は入力バウンドが多数であるため、データのワーカー数が多くなります。そのため、ワーカー スレッド数を維持するのに十分な RAM がない場合は、トレーニングが非常に遅くなってしまいます。最適な結果を得るには、n1-highmem-96 マシンタイプを選択してください。

ブートディスクとして Debian GNU/Linux 9 Stretch + PyTorch/XLA を選択します。ImageNet をダウンロードする場合は、300 GB 以上のディスクサイズを指定します。架空データのみを使用する場合は、デフォルトのディスクサイズ(20 GB)を指定します。

このチュートリアルでは、両方のデータセットの使用を提案します。

TPU を作成する

  1. コンソールに移動して TPU を作成します。
  2. [名前] で、TPU Pod の名前を指定します。
  3. [ゾーン] で、Cloud TPU に使用するゾーンを指定します。VM と同じゾーンにあることを確認します。
  4. [TPU タイプ] で、Cloud TPU タイプを選択します。最適な結果を得るには、v3-8 TPU を選択してください。
  5. [TPU ソフトウェア バージョン] で、最新の安定リリース(pytorch-0.5)を選択します。
  6. デフォルトのネットワークを使用する。
  7. IP アドレス範囲を設定します(10.240.0.0 など)。

架空のデータセットを使用してトレーニングする

初めて実行する際は、実際の ImageNet セットではなく、架空のデータセットを使用することをおすすめします。fake_data が自動的に VM にインストールされ、処理にかかる時間とリソースを節約できます。

トレーニングには conda または Docker を使用できます。Docker に不慣れな場合は、conda を使用してください。

  • Debian GNU/Linux 9 Stretch + PyTorch/XLA ディスク内の conda 環境でパッケージ化されたコードのバージョンを実行します。
  • Docker を使用して実行します。最新の docker イメージを pull して、トレーニングを開始できるまでに数分かかります。

オプション 1

# Fill in your the name of your VM and the zone.
$ gcloud beta compute  ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(vm)$ conda activate torch-xla-0.5
(torch-xla-0.5)$  python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --fake_data --model=resnet50 --num_epochs=10 --num_workers=64 --batch_size=128 --log_steps=20

オプション 2

# Fill in your the name of your VM and the zone.
$ gcloud beta compute ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ docker run --shm-size 16G -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" gcr.io/tpu-pytorch/xla:r0.5 python3 pytorch/xla/test/test_train_imagenet.py --model=resnet50 --num_epochs=10 --num_workers=64 --fake_data

実際のデータセットを使用してトレーニングする

--fake_data フラグを付けて実行して何も問題がなさそうであれば、ImageNet などの実データのトレーニングを試すことができます。

一般的に、test_train_imagenet.py では torchvision.datasets.ImageFolder が使用されているので、適切に構造化された任意のデータセットを使用できます。ImageFolder のドキュメントをご覧ください。

データセットが ~/imagenet に保存されているという前提で、実際のデータの変更を次のように提案されることがあります。

オプション 1

# Fill in your the name of your VM and the zone.
$ gcloud beta compute  ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(vm)$ ulimit -n 10240
(vm)$ conda activate torch-xla-0.5
(torch-xla-0.5)$ python /usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py --datadir=~/imagenet --model=resnet50 --num_epochs=90 --num_workers=64 --batch_size=128 --log_steps=200

オプション 2

# Fill in your the name of your VM and the zone.
$ gcloud beta compute ssh "your-VM-name" --zone "your-zone".
(vm)$ export TPU_IP_ADDRESS=your-ip-address
(vm)$ docker run --shm-size 128G -v ~/imagenet:/tmp/imagenet -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" gcr.io/tpu-pytorch/xla:r0.5 python3 pytorch/xla/test/test_train_imagenet.py --model=resnet50 --num_epochs=90 --num_workers=64 --log_steps=200 --datadir=/tmp/imagenet

n1-highmem-96 VM とv3-8 TPU を使用する ImageNet データセット全体では、通常、最初のエポックには 20 分、その後のエポックには約 11 分かかります。モデルは、90 エポックの時点で最上位の約 76% の精度に到達します。

追加の CPU 割り当てをリクエストする

リソースの追加は、数日前に計画してリクエストし、リクエストが時間的な余裕をもって実行されるようにします。

  1. [割り当て] ページに移動します。

    [割り当て] ページに移動

  2. [サービス] メニューから [Cloud TPU API] を選択します。
  3. CPU を使用するリージョンまたはゾーンを選択します。
  4. [指標] メニューから [なし] を選択し、検索ボックスに CPU を入力します。
  5. [CPU] を選択します。
  6. [Compute Engine API] - [CPU] を選択し、ページの上部にある[割り当てを編集]をクリックします。
  7. リクエストしている割り当て量と説明を入力し、[Done] をクリックします。サービス プロバイダに承認を求めるリクエストが送信されます。

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud Platform アカウントに課金されないようにする手順は次のとおりです。

  1. Compute Engine VM を終了します。
  2. VM を削除します。
  3. TPU を削除します。

次のステップ

次のように PyTorch colabs を試す