Cloud TPU での PyTorch モデルのトレーニング

このドキュメントでは、Cloud TPU で PyTorch モデルを実行する方法について説明します。Cloud TPU で PyTorch モデルをトレーニングする方法は、使用している Cloud TPU アーキテクチャによって異なります。詳細については、システム アーキテクチャをご覧ください。

Cloud TPU VM

Cloud TPU VM にはユーザー VM がなく、モデルとコードは TPU ホストマシン上の VM で実行されます。TPU ホストには、直接 SSH で接続できます。任意のコード インストール パッケージの実行、ログの表示、コードのデバッグなどを行えます。

PyTorch/XLA は、新しい PyTorch/XLA プログラムが実行されるたびに TensorFlow ローカル サーバーを作成します。XRT（XLA Runtime）クライアントは、ローカルの TensorFlow サーバーに接続します。

画像

TPU VM 上で PyTorch モデルをトレーニングすることの詳細については、Cloud TPU VM PyTorch クイックスタートをご覧ください。

以前の TPU ノード

レガシー TPU ノードを使用する場合は、TPU ごとにユーザー VM が必要です。ユーザー VM は gRPC 呼び出しで TPU ノードと通信します。TPU ノードは TensorFlow サーバーを実行します。XRT（XLA ランタイム）クライアントは、ネットワークを介して TensorFlow サーバーに接続します。

画像