この古いバージョンの AI Platform Training は非推奨になりました。2025 年 1 月 31 日を過ぎると Google Cloud で使用できなくなります。リソースを Vertex AI カスタムトレーニングに移行することで、AI Platform にはない新しい機械学習機能を利用できます。

TPU を使用したモデルのトレーニング

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=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
```

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 についてのトラブルシューティングとよくある質問を参照して、診断と問題解決を行う。

GPU の使用

TPU を使用したモデルのトレーニング

Google Cloud 環境の設定

Cloud TPU にプロジェクトへのアクセスを許可する

Console

gcloud

例: サンプル MNIST モデルのトレーニング

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

TPU を提供するリージョンの指定

TensorFlow と AI Platform Training のバージョニング

TPU gRPC サーバーへの接続

TensorFlow コードでの TPU の使用

PyTorch コードでの TPU の使用

カスタム TPU マシンの構成

基本的な TPU 対応マシン

従来のマシンタイプ構成の TPU ワーカー

Compute Engine マシンタイプ構成の TPU ワーカー

TPU v2

TPU v3（ベータ版）

TPU Pod の使用

TPU v2 Pod

TPU v3 Pod

TPU ワーカーでのビルド済み PyTorch コンテナの使用

TPU ワーカーでのカスタム コンテナの使用

TPU v2

TPU v3（ベータ版）

TPU のプロビジョニング後の TPUClusterResolver の使用

次のステップ

TPU ワーカーでのカスタムコンテナの使用

TPU のプロビジョニング後の `TPUClusterResolver` の使用