Cloud TPU 上のソフトウェア バージョンの切り替え

概要

TPU で実行されるフレームワークのソフトウェア バージョンは、ローカル VM で実行されているバージョンと一致する必要があります。このソフトウェア バージョンは、TPU を削除して再作成せずに、実行中の Cloud TPU で切り替えられるようになりました。そのため、特定の夜間バージョンのソフトウェア フレームワークで Cloud TPU を構成することもできます。これらのフレームワークのサポートされているバージョンを選択することをおすすめします。

用途

バージョンを切り替えるには、cloud-tpu-client Python ライブラリを使用することをおすすめします。

TensorFlow の使用例

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

これにより、ローカル VM で実行される TensorFlow バージョンと一致するように Cloud TPU が構成されます。これには、公式リリースと、日付の付いた夜間ビルドが含まれます。

configure_tpu_version API の restart_type パラメータは、バージョンの切り替え時の TPU 再起動動作を定義します。オプションは「常時」(デフォルト)で、「ifNeeded」です。

  • 「always」は、ステータス UNHEALTHY_TENSORFLOW などの TPU を修正する場合や、前回の実行でリソースが漏洩してメモリ不足(OOM)エラーを返す場合などに使用できます。このオプションを設定した場合、新しいフレームワーク バージョンがインストールされていない場合でも TPU は再起動されます。

  • 「ifNeeded」は、すでに適切なバージョンになっている場合はランタイムが再起動されないため、トレーニング スクリプトにかなりのスタートアップ時間を追加する必要がないため、便利です。このオプションを設定した場合、TPU は正しいフレームワーク バージョンがインストールされていない場合にのみ再起動されます。

このライブラリは Cloud TPU と直接通信します。そのため、このコードは同じネットワーク内の VM で実行する必要があります。残りのモデルのコード内で実行することをおすすめします。

その他のソフトウェア オプション

TensorFlow には、正しいバージョンを構成する最も簡単な方法である tf.__version__ 文字列が含まれています。その他のソフトウェア オプションは次のとおりです。

  • PyTorch - pytorch-0.1pytorch-0.5-dev20200320pytorch-nightly
  • Jax - tpu_drivertpu_driver0.1-dev20200320tpu_driver_nightly

たとえば、PyTorch の最新の夜間ビルドで実行するように TPU を構成します。

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()