Tensor Processing Unit(TPU)は、機械学習ワークロードを高速化するために使用される、Google 独自開発の ASIC です。Cloud TPU を使用して、AI Platform Training でトレーニング ジョブを実行できます。AI Platform Training にはジョブ管理インターフェースが用意されているので、TPU を自分で管理する必要はありません。CPU や GPU でトレーニングを実行する場合と同じように AI Platform Training jobs
API を使用できます。
高レベルの TensorFlow API は、Cloud TPU ハードウェア上でモデルを実行する場合に役立ちます。
Google Cloud 環境の設定
スタートガイドのセットアップ セクションに沿って、Google Cloud 環境を構成します。
Cloud TPU にプロジェクトへのアクセスを許可する
Google Cloud プロジェクトに関連付けられた Cloud TPU サービス アカウント名を承認する手順は次のとおりです。
projects.getConfig
を呼び出して、Cloud TPU サービス アカウント名を取得します。例:PROJECT_ID=PROJECT_ID curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
API から返された
serviceAccountProject
とtpuServiceAccount
フィールドの値を保存します。Cloud TPU サービス アカウントを初期化します。
curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ -H "Content-Type: application/json" -d '{}' \ https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
Cloud ML サービス エージェントのロールを使用して、Cloud TPU サービス アカウントをプロジェクト内のメンバーとして追加します。Google Cloud Console または gcloud
コマンドを使用して、次の手順を行います。
Console
- Google Cloud Console にログインし、TPU を使用しているプロジェクトを選択します。
- [IAM と管理] > [IAM] を選択します。
- [追加] ボタンをクリックして、プロジェクトにメンバーを追加します。
- [メンバー] テキスト ボックスに TPU サービス アカウントを入力します。
- [ロール] プルダウン リストをクリックします。
- Cloud ML サービス エージェントのロールを有効にします([サービス エージェント] > [Cloud ML サービス エージェント])。
gcloud
プロジェクト ID と Cloud TPU サービス アカウントが格納されている環境変数を設定します。
PROJECT_ID=PROJECT_ID SVC_ACCOUNT=[email protected]
Cloud TPU サービス アカウントに
ml.serviceAgent
ロールを付与します。gcloud projects add-iam-policy-binding $PROJECT_ID \ --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
サービス アカウントへのロールの付与については、IAM のドキュメントをご覧ください。
例: サンプル MNIST モデルのトレーニング
このセクションでは、TPU とランタイム バージョン 2.11 を使用してサンプル MNIST モデルをトレーニングする方法について説明します。この例のジョブでは、マシン構成にあらかじめ定義された BASIC_TPU
スケールティアを使用します。ガイドのこの後のセクションでは、カスタム構成を設定する方法について説明します。
この例では、gcloud CLI がインストールされた Bash シェルを使用していることを前提としています。次のコマンドを実行してコードを取得し、トレーニング ジョブを AI Platform Training に送信します。
TensorFlow のリファレンス モデルのコードをダウンロードし、サンプルコードのあるディレクトリに移動します。
git clone https://github.com/tensorflow/models.git \ --branch=v2.11.0 \ --depth=1 cd models
models
ディレクトリにsetup.py
ファイルを作成します。これにより、gcloud ai-platform jobs submit training
コマンドは、トレーニング コードの tarball を作成するときに、必要なすべてのサブパッケージをmodels/official
ディレクトリに含め、トレーニング ジョブ実行時に AI Platform Training が TensorFlow データセットを依存関係としてインストールします。このトレーニング コードは、MNIST データを読み込むために TensorFlow データセットに依存します。setup.py
ファイルを作成するには、シェルで次のコマンドを実行します。cat << END > setup.py from setuptools import find_packages from setuptools import setup setup( name='official', install_requires=[ 'tensorflow-datasets~=3.1', 'tensorflow-model-optimization>=0.4.1' ], packages=find_packages() ) END
gcloud ai-platform jobs submit training
コマンドを使用して、トレーニング ジョブを送信します。gcloud ai-platform jobs submit training tpu_mnist_1 \ --staging-bucket=gs://BUCKET_NAME \ --package-path=official \ --module-name=official.vision.image_classification.mnist_main \ --runtime-version=2.11 \ --python-version=3.7 \ --scale-tier=BASIC_TPU \ --region=us-central1 \ -- \ --distribution_strategy=tpu \ --data_dir=gs://tfds-data/datasets \ --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
BUCKET_NAME は、Google Cloud プロジェクト内の Cloud Storage バケットの名前に置き換えます。gcloud CLI はパッケージ化されたトレーニング コードをこのバケットにアップロードし、AI Platform Training はバケットにトレーニング出力を保存します。
トレーニング ジョブをモニタリングします。ジョブが完了したら、
gs://BUCKET_NAME/tpu_mnist_1_output
ディレクトリでその出力を表示できます。
Cloud TPU でのモデルのトレーニングの詳細
このセクションでは、AI Platform Training 上で Cloud TPU を使用してジョブの構成とモデルのトレーニングを行う方法について詳しく説明します。
TPU を提供するリージョンの指定
TPU が利用可能なリージョンでジョブを実行する必要があります。次のリージョンで現在 TPU にアクセスできるようになっています。
us-central1
europe-west4
モデル トレーニングやオンライン / バッチ予測などの AI Platform Training サービスで利用可能なリージョンの詳細については、リージョンのガイドをご覧ください。
TensorFlow と AI Platform Training のバージョニング
AI Platform Training ランタイム バージョン 1.15、2.1、2.2、2.3、2.4、2.5、2.6、2.7、2.8、2.9、2.11 は Cloud TPU でのモデルのトレーニングに使用できます。AI Platform Training ランタイム バージョンと対応する TensorFlow バージョンに関する詳細をご覧ください。
バージョニング ポリシーは Cloud TPU と同じです。トレーニングのジョブ リクエストでは、TPU で使用可能なランタイム バージョンを指定し、トレーニング コードで使用されている TensorFlow バージョンと一致することを確認してください。
TPU gRPC サーバーへの接続
TensorFlow プログラムでは、TPUClusterResolver
を使用して TPU VM 上で実行されている TPU gRPC サーバーに接続します。
TPU の使用に関する TensorFlow ガイドでは、TPUStrategy
配信戦略で TPUClusterResolver
を使用する方法を説明しています。
ただし、AI Platform Training で実行するコードに TPUClusterResolver
を使用する場合は、重要な変更点があります。TPUClusterResolver
インスタンスの作成時に引数を指定しないでください。tpu
、zone
、project
の各キーワード引数がすべてデフォルト値の None
に設定されている場合、AI Platform Training は環境変数を通じて必要な接続の詳細を含むクラスタ リゾルバを自動的に提供します。
次の TensorFlow 2 の例は、AI Platform Training でトレーニングするためのクラスタ リゾルバと分散戦略を初期化する方法を示しています。
import tensorflow as tf
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)
TensorFlow コードでの TPU の使用
マシンで TPU を使用するには、TensorFlow 2 の TPUStrategy
API を使用します。これを行う方法については、TPU の使用に関する TensorFlow ガイドをご覧ください。
TensorFlow 1 で TPU を使用してトレーニングするには、代わりに TPUEstimator
API を使用します。これを行う方法については、TPUEstimator
API に関する Cloud TPU ガイドをご覧ください。
Cloud TPU のドキュメントには、Cloud TPU で利用できる簡単な TensorFlow 演算のリストも記載されています。
PyTorch コードでの TPU の使用
ビルド済みの PyTorch コンテナを使用する際に TPU を使用するには、torch_xla
パッケージを使用します。PyTorch ドキュメントのトレーニングで TPU に torch_xla
を使用する方法について学習します。torch_xla
の他の使用例については、PyTorch XLA GitHub リポジトリのチュートリアルをご覧ください。
AI Platform Training で TPU を使用してトレーニングする場合は、複数の XLA デバイスではなく、単一の XLA デバイスを使用します。
このページの次のセクションで、PyTorch と TPU のトレーニング ジョブの構成もご覧ください。
カスタム TPU マシンの構成
TPU トレーニング ジョブは、2 つの VM による構成で実行されます。一方の VM(マスター)が Python コードを実行します。マスターは、TPU ワーカーで実行されている TensorFlow サーバーを駆動します。
AI Platform Training で TPU を使用するには、次の 3 つの方法のいずれかで TPU 対応マシンにアクセスするようにトレーニング ジョブを構成します。
BASIC_TPU
スケール階層を使用する。この方法で TPU v2 アクセラレータにアクセスできます。- マスター VM で
cloud_tpu
ワーカーと従来のマシンタイプを使用する。この方法で TPU v2 アクセラレータにアクセスできます。 - マスター VM で
cloud_tpu
ワーカーと Compute Engine マシンタイプを使用する。この方法で TPU v2 アクセラレータまたは TPU v3 アクセラレータにアクセスできます。TPU v3 アクセラレータはベータ版で利用できます。
基本的な TPU 対応マシン
スケール階層を BASIC_TPU
に設定して、前述のサンプルを実行した場合と同様に、マスター VM と、8 つの TPU v2 コアを搭載した 1 つの TPU を使用している TPU VM を取得します。
従来のマシンタイプ構成の TPU ワーカー
また、マスター VM でさらに多くのコンピューティング リソースが必要な場合は、次のようにカスタムマシン構成を設定することもできます。
- スケール階層を
CUSTOM
に設定する。 - ジョブの要件に合った従来のマシンタイプを使用するようにマスター VM を構成する。
workerType
をcloud_tpu
に設定して、8 つの TPU v2 コアを搭載した 1 つの Cloud TPU を使用している TPU VM を取得します。workerCount
を 1 に設定する。- Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。
parameterServerCount
がゼロより大きい場合、サービスはジョブ リクエストを拒否します。
次の例は、このタイプの構成を使用する config.yaml
ファイルを示しています。
trainingInput:
scaleTier: CUSTOM
masterType: complex_model_m
workerType: cloud_tpu
workerCount: 1
Compute Engine マシンタイプ構成の TPU ワーカー
マスター VM の Compute Engine マシンタイプにカスタムマシン構成を設定し、TPU VM に接続された acceleratorConfig
を設定することもできます。
このタイプの構成を使用して、8 つの TPU v2 コアを搭載した TPU ワーカー(acceleratorConfig
なしの構成に類似)または 8 つの v3 コアを搭載した TPU ワーカー(ベータ版)を設定できます。TPU v2 アクセラレータと TPU v3 アクセラレータの違いをご確認ください。
Compute Engine マシンタイプを使用すると、次のようにマスター VM をより柔軟に構成できます。
- スケール階層を
CUSTOM
に設定する。 - ジョブ要件に合った Compute Engine マシンタイプを使用するようにマスター VM を構成する。
workerType
をcloud_tpu
に設定する。acceleratorConfig
フィールドを含むworkerConfig
を追加します。そのacceleratorConfig
内で、type
をTPU_V2
またはTPU_V3
に、count
を8
に設定します。それ以外の数の TPU コアを接続することはできません。workerCount
を 1 に設定する。- Cloud TPU を使用する場合は、パラメータ サーバーを指定しない。
parameterServerCount
がゼロより大きい場合、サービスはジョブ リクエストを拒否します。
次の例は、このタイプの構成を使用する config.yaml
ファイルを示しています。
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3(ベータ版)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3
count: 8
TPU Pod の使用
TPU Pod は、専用の高速ネットワーク インターフェースによって接続された TPU デバイスのコレクションです。1 つの TPU Pod に最大 2,048 個の TPU コアを設定して、複数の TPU で処理負荷を分散できます。
TPU Pod を使用するには、まず割り当ての増加リクエストを送信する必要があります。
次の config.yaml
ファイルの例は、TPU Pod の使用方法を示しています。
TPU v2 Pod
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2_POD
count: 128
TPU v3 Pod
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3_POD
count: 32
TPU タイプごとに使用できる Pod コアの数には制限があります。有効な構成:
TPU Pod タイプ | 使用可能な Pod のコア数 |
---|---|
TPU_V2_POD |
32, 128, 256, 512 |
TPU_V3_POD |
32, 128, 256 |
TPU Pod コアを最大限に活用する方法の詳細については、TPU Pod に関する Cloud TPU のドキュメントをご覧ください。
TPU ワーカーでのビルド済み PyTorch コンテナの使用
TPU を使用して PyTorch トレーニングを実行する場合、トレーニング ジョブの trainingInput
に tpuTfVersion
フィールドを指定する必要があります。トレーニングに使用するビルド済みの PyTorch コンテナのバージョンに合わせて tpuTfVersion
を設定します。
AI Platform Training は、次のビルド済み PyTorch コンテナの TPU によるトレーニングをサポートしています。
コンテナ イメージの URI | tpuTfVersion |
---|---|
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 |
pytorch-1.11 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 |
pytorch-1.10 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 |
pytorch-1.9 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 |
pytorch-1.7 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 |
pytorch-1.6 |
たとえば、PyTorch 1.11 のビルド済みコンテナを使用してトレーニングするには、次の config.yaml
ファイルを使用してトレーニングを構成します。
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
tpuTfVersion: pytorch-1.11
acceleratorConfig:
type: TPU_V2
count: 8
このページの前のセクションの PyTorch コードでの TPU の使用もご覧ください。
TPU ワーカーでのカスタム コンテナの使用
TPU をサポートする AI Platform Training ランタイム バージョンのいずれかを使用せずに、TPU ワーカーでカスタム コンテナを実行する場合は、トレーニング ジョブの送信時に追加の構成フィールドを指定する必要があります。tpuTfVersion
を、コンテナで使用する TensorFlow のバージョンが含まれているランタイム バージョンに設定します。TPU でのトレーニングのために現在サポートされているランタイム バージョンを指定する必要があります。
カスタム コンテナを使用するジョブを構成しているので、AI Platform のトレーニングではトレーニング ジョブの実行時にこのランタイム バージョンの環境を使用しません。ただし、カスタム コンテナで使用する TensorFlow のバージョン用に TPU ワーカーを適切に準備できるように、AI Platform のトレーニングでこのフィールドが必要です。
次の例では、前のセクションと同様の TPU 構成の config.yaml
ファイルを示しますが、この場合はマスター VM と TPU ワーカーが異なるカスタム ワーカーで動作します。
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3(ベータ版)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V3
count: 8
トレーニング ジョブを送信するために gcloud beta ai-platform jobs submit training
コマンドを使用する場合は、config.yaml
ファイルではなく --tpu-tf-version
フラグを使用して tpuTfVersion
API フィールドを指定できます。
TPU のプロビジョニング後の TPUClusterResolver
の使用
カスタム コンテナを使用する場合は、TPUClusterResolver
を呼び出して使用する前に TPU がプロビジョニングされるのを待つ必要があります。次のサンプルコードは、TPUClusterResolver
ロジックの処理方法を示しています。
def wait_for_tpu_cluster_resolver_ready():
"""Waits for `TPUClusterResolver` to be ready and return it.
Returns:
A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
return None.
Raises:
RuntimeError: if failed to schedule TPU.
"""
tpu_config_env = os.environ.get('TPU_CONFIG')
if not tpu_config_env:
tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
return None
tpu_node = json.loads(tpu_config_env)
tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)
num_retries = 40
for i in range(num_retries):
try:
tpu_cluster_resolver = (
tf.contrib.cluster_resolver.TPUClusterResolver(
tpu=[tpu_node['tpu_node_name']],
zone=tpu_node['zone'],
project=tpu_node['project'],
job_name='worker'))
tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
if 'worker' in tpu_cluster_resolver_dict:
tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
return tpu_cluster_resolver
except Exception as e:
if i < num_retries - 1:
tf.logging.info('Still waiting for provisioning of TPU VM instance.')
else:
# Preserves the traceback.
raise RuntimeError('Failed to schedule TPU: {}'.format(e))
time.sleep(10)
# Raise error when failed to get TPUClusterResolver after retry.
raise RuntimeError('Failed to schedule TPU.')
詳細については、カスタム コンテナを使用した分散トレーニングをご覧ください。
次のステップ
- AI Platform Training でのトレーニング モデルについて詳しく学ぶ。
- AI Platform Training のハイパーパラメータ調整について学ぶ。特に、Cloud TPU を使用したハイパーパラメータ調整を詳しく学びます。
- Cloud TPU のリファレンス モデルをさらに確認する。
- Cloud TPU のベスト プラクティスに従って Cloud TPU のモデルを最適化する。
- Cloud TPU についてのトラブルシューティングとよくある質問を参照して、診断と問題解決を行う。