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 ランタイム)クライアントは、ローカルの TensorFlow サーバーに接続します。

画像

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

以前の TPU ノード

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

画像