PyTorch はついに Google Cloud TPU を使用できるようになりました。これにより推論パフォーマンスが 4 倍向上します。これをどのように活用できるでしょうか?

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

Facebook は PyTorch 開発者会議で PyTorch 1.3 を正式にリリースし、Colab でも呼び出せる Google Cloud TPU の完全なサポートを発表しました。

機械学習開発者はこれまでも Colab で PyTorch を使用できましたが、クラウド TPU をサポートするのは今回が初めてです。これは、高価な GPU を購入する必要がなく、クラウドで独自のモデルをトレーニングできることも意味します。

Google Cloud Platform の新規ユーザーの場合は、300 ドル分の無料クレジットを獲得できます。

現在、PyTorch は Github でサンプルコードを公式に提供しており、Google Cloud TPU を使用してモデルを無料でトレーニングし、Colab で推論を実行する方法を説明しています。

ResNet-50のトレーニング

PyTorch は、クラウド TPU デバイス上で ResNet-50 モデルをトレーニングするケースを初めて導入しました。 Cloud TPU を使用して他の画像分類モデルをトレーニングする場合も、操作は同様です。

トレーニングの前に、コンソールに移動して、VM の名前とリージョンを指定して新しい VM インスタンスを作成する必要があります。

実際のデータで Resnet50 をトレーニングする場合は、CPU の数が最も多いマシンタイプを選択する必要があります。最良の結果を得るには、 n1-highmem-96マシンタイプを選択してください。

次に、Debian GNU/Linux 9 Stretch + PyTorch/XLA ブートディスクを選択します。実際の ImageNet データを使用してトレーニングする予定の場合は、少なくとも 300 GB のディスク容量が必要になります。トレーニングに偽のデータを使用する場合、デフォルトのディスクサイズは 20 GB のみです。

TPUの作成

コンソールに移動して TPU を作成します。
「名前」に、TPU ポッドの名前を指定します。
[ゾーン] で、Cloud TPU のゾーンを指定します。その際、前に作成した VM と同じゾーンであることを確認します。
「TPU タイプ」で TPU タイプを選択します。最良の結果を得るには、 v3-8 TPU (8 v3) を選択します。
「TPU ソフトウェアバージョン」で、最新の安定バージョンを選択します。
デフォルトのネットワークを使用します。
IP アドレスの範囲を設定します (例: 10.240.0.0)。

fake_data は仮想マシンに自動的にインストールされ、必要な時間とリソースが少なくなるため、最初の実行時にはトレーニングに偽のデータを使用することをお勧めします。トレーニングには conda または Docker を使用できます。

fake_data でのテストに成功したら、ImageNet などの実際のデータでのトレーニングを開始できます。

conda を使用したトレーニング:

 # VM の名前とゾーンを入力します。
 $ gcloud beta compute ssh "your-VM-name" --zone "your-zone" 。
 (vm)$ export TPU_IP_ADDRESS=IPアドレス
(vm)$ エクスポート XRT_TPU_CONFIG= "tpu_worker;0;$TPU_IP_ADDRESS:8470"  
 (vm)$ ulimit -n 10240  
 (vm)$ conda をアクティブ化 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

Docker を使用したトレーニング:

 # VM の名前とゾーンを入力します。
 $ gcloud beta compute ssh "your-VM-name" --zone "your-zone" 。
 (vm)$ export TPU_IP_ADDRESS=IPアドレス
(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 を使用してトレーニングする場合、通常、最初のエポックでは約 20 分、後続のエポックでは約 11 分かかります。このモデルは、90 エポック後に約 76% のトップ 1 精度を達成します。

Google Cloud による後続の課金を回避するには、トレーニングが完了したら仮想マシンと TPU を削除することを忘れないでください。

パフォーマンスはGPUの4倍

トレーニングが完了したら、Colab にモデルをインポートできます。

ノートブックファイルを開き、メニューバーの[ランタイム] で[ランタイムタイプの変更]を選択し、ハードウェアアクセラレータのタイプを TPU に変更します。

まず次のコードセルを実行して、Colab 上の TPU にアクセスできることを確認します。

インポートOS
 assert os.environ['COLAB_TPU_ADDR'], '編集 > ノートブック設定 > ハードウェアアクセラレータから TPU を選択してください'

次に、互換性のある PyTorch/TPU コンポーネントを Colab にインストールします。

 DIST_BUCKET = "gs://tpu-pytorch/wheels"  
 TORCH_WHEEL = "torch-1.15-cp36-cp36m-linux_x86_64.whl"  
 TORCH_XLA_WHEEL = "torch_xla-1.15-cp36-cp36m-linux_x86_64.whl"  
 TORCHVISION_WHEEL = "torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl"  
 # Colab TPU互換PyTorch/TPUホイールと依存関係をインストールする
!pip アンインストール -y torch torchvision
 !gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" 。
 !gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" 。
 !gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" 。
 !pip で"$TORCH_WHEEL"をインストールします 
 !pip で"$TORCH_XLA_WHEEL"をインストールします 
 !pip で"$TORCHVISION_WHEEL"をインストールします 
 !sudo apt-get libomp5 をインストール