Cloud TPU（TF 2.x）でのトレーニング

このチュートリアルでは、tf.distribute.TPUStrategy を使用して Cloud TPU で Keras ResNet-RS モデルをトレーニングする方法を説明します。ResNet-RS の詳細については、ResNets の再確認: トレーニングとスケーリング戦略の改善をご覧ください。

Cloud TPU に慣れていない場合は、TPU と Compute Engine VM を作成する方法について、クイックスタートを参照することを強くおすすめいたします。

目標

  • データセットとモデルの出力を格納する Cloud Storage バケットを作成します。
  • ImageNet データセットに類似したフェイク ImageNet データセットを準備します。
  • トレーニング ジョブを実行します。
  • 出力結果を確認します。

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

  • Compute Engine
  • Cloud TPU
  • Cloud Storage

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

  1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

  2. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  3. Google Cloud プロジェクトで課金が有効になっていることを確認します。

  4. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

    プロジェクト セレクタに移動

  5. Google Cloud プロジェクトで課金が有効になっていることを確認します。

    6. このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

リソースを設定する

このセクションでは、チュートリアルで使用する Cloud Storage のバケット、VM、Cloud TPU の各リソースを設定する方法を説明します。

  1. Cloud Shell ウィンドウを開きます。

    Cloud Shell を開く

  2. プロジェクト ID の変数を作成します。

    export PROJECT_ID=project-id

  3. Cloud TPU を作成するプロジェクトを使用するように Google Cloud CLI を構成します。

    gcloud config set project ${PROJECT_ID}

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。

  4. Cloud TPU プロジェクトのサービス アカウントを作成します。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

    このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。 

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com

  5. 次のコマンドを使用して Cloud Storage バケットを作成します。

    gsutil mb -p ${PROJECT_ID} -c standard -l europe-west4 gs://bucket-name

    この Cloud Storage バケットには、モデルのトレーニングに使用するデータとトレーニング結果が格納されます。このチュートリアルで使用する gcloud compute tpus execution-groups ツールは、前の手順で設定した Cloud TPU サービス アカウントのデフォルトの権限を設定します。権限の詳細な設定が必要な場合は、アクセスレベル権限をご覧ください。

    バケットのロケーションは、Compute Engine（VM）および Cloud TPU ノードと同じリージョンにする必要があります。

  6. gcloud コマンドを使用して、Compute Engine VM を起動します。

    $ gcloud compute tpus execution-groups create \
 --vm-only \
 --name=resnet-rs-tutorial \
 --zone=europe-west4-a \
 --disk-size=300 \
 --machine-type=n1-standard-16 \
 --tf-version=2.12.0

    コマンドフラグの説明

    vm-only
    VM のみを作成し、TPU は作成しないでください。
    name
    作成する TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    disk-size
    gcloud compute tpus execution-groups コマンドで作成された VM のハードディスクのサイズ（GB）。
    machine-type
    作成する Compute Engine VM のマシンタイプ。

    gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

  7. プロンプトが表示されたら、y キーを押して Cloud TPU リソースを作成します。

    gcloud compute tpus execution-groups コマンドの実行が終了したら、shell プロンプトが username@projectname から username@vm-name に変更されたことを確認します。 変更されていれば、Compute Engine VM にログインしていることになります。

    gcloud compute ssh resnet-rs-tutorial --zone=europe-west4-a

    これらの手順を続行する場合は、Compute Engine インスタンスで (vm)$ で始まる各コマンドを実行します。

  8. 必要なパッケージをインストールします。

    $ pip3 install tensorflow-text==2.8.1 --no-deps

Cloud Storage バケットの変数を設定する

次の環境変数を設定します。bucket-name を Cloud Storage バケットの名前に置き換えます。

(vm)$ export STORAGE_BUCKET=gs://bucket-name

(vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-rs-2x
(vm)$ export IMAGENET_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
(vm)$ export PYTHONPATH=/usr/share/models
(vm)$ export TPU_NAME=resnet-rs-tutorial

トレーニング アプリケーションでは、Cloud Storage でトレーニング データにアクセスできる必要があります。また、トレーニング アプリケーションは、Cloud Storage バケットを使用してトレーニング中にチェックポイントを保存します。

Cloud TPU 単一デバイスのトレーニングと評価

ImageNet は画像データベースです。このデータベース内では画像が階層に編成されていて、階層の各ノードを数百、数千もの画像で表しています。

このチュートリアルでは、ImageNet の完全版のデータセットの fake_imagenet と呼ばれるデモバージョンを使用しています。このデモバージョンを使用すると、ストレージ容量と所要時間を ImageNet の完全版のデータセットに対してモデルを実行する際に通常必要となるものより抑えながらチュートリアルを試すことができます。

fake_imagenet データセットは Cloud Storage の次のロケーションにあります。

gs://cloud-tpu-test-datasets/fake_imagenet

fake_imagenet データセットは、Cloud TPU の使用方法を理解し、エンドツーエンドのパフォーマンスを検証する場合にのみ役立ちます。精度の数値と保存されたモデルは意味がありません。

ImageNet データセット全体をダウンロードして処理する方法については、ImageNet データセットのダウンロード、前処理、アップロードをご覧ください。

  1. gcloud コマンドを使用して Cloud TPU を作成します。

    $ gcloud compute tpus execution-groups create \
 --tpu-only \
 --accelerator-type=v3-8  \
 --name=resnet-rs-tutorial \
 --zone=europe-west4-a \
 --tf-version=2.12.0

    コマンドフラグの説明

    tpu-only
    Cloud TPU のみを作成し、VM は作成しないでください。
    accelerator-type
    作成する Cloud TPU のタイプ。
    name
    作成する TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    tf-version
    VM にインストールされている Tensorflow のバージョン。

  2. TPU_NAME 名前変数を設定します。

    (vm)$ export TPU_NAME=resnet-rs-tutorial

  3. トレーニング スクリプトを実行します。

    (vm)$ python3 /usr/share/models/official/vision/beta/train.py \
--experiment=resnet_rs_imagenet \
--mode=train_and_eval \
--model_dir=$MODEL_DIR \
--tpu=$TPU_NAME \
--config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
--params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*, trainer.train_steps=100"

    コマンドフラグの説明

    experiment
    実行するテストの名前。
    mode
    スクリプトを実行するモード。有効な値は、「train」、「eval」、または「train_and_eval」です。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    tpu
    使用する TPU の名前。
    config_file
    スクリプト構成ファイルへのパス。
    params_override
    スクリプト構成ファイルの設定をオーバーライドします。

これにより、ResNet-RS を 100 トレーニング ステップでトレーニングし、v3-8 TPU ノードでは 5 分以内に完了します。トレーニング スクリプトの出力には、次のようなテキストが含まれます。

{
  'train_loss': 1.435225,
  'train_accuracy': 0.00084427913
}

このトレーニング スクリプトも評価を実行します。評価出力には次のようなテキストが含まれます。

Run stats:
{
  'eval_loss': 0.861013,
  'eval_acc': 0.001,
  'train_loss': 1.435225,
  'train_acc': 0.00084427913,
  'step_timestamp_log': [
    'BatchTimestamp<batch_index: 0,
    timestamp: 1606330585.7613473>',
    'BatchTimestamp<batch_index: 500,
    timestamp: 1606330883.8486104>',
    'BatchTimestamp<batch_index: 1000,
    timestamp: 1606331119.515312>',
    'BatchTimestamp<batch_index: 1251,
    timestamp: 1606331240.7516596>'
  ],
  'train_finish_time': 1606331296.395158,
  'avg_exp_per_second': 1951.6983246161021
}

ResNet-RS モデルを収束するようにトレーニングするには、次のスクリプトに示すように trainer.train_steps=100 引数を省略します。トレーニングと評価は一緒に行われます。

(vm)$ python3 /usr/share/models/official/vision/beta/train.py \
  --experiment=resnet_rs_imagenet \
  --mode=train_and_eval \
  --model_dir=$MODEL_DIR \
  --tpu=$TPU_NAME \
  --config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
  --params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*"

コマンドフラグの説明

experiment
実行するテストの名前。
mode
スクリプトを実行するモード。有効な値は、「train」、「eval」、または「train_and_eval」です。
model_dir
トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
tpu
使用する TPU の名前。
config_file
スクリプト構成ファイルへのパス。
params_override
スクリプト構成ファイルの設定をオーバーライドします。

トレーニングと評価は fake_imagenet データセットに対して行われているため、出力結果には、実際のデータセットでトレーニングと評価を行った場合の出力は反映されません。

この時点で、このチュートリアルを終了して、Google Cloud リソースをクリーンアップすることも、Cloud TPU Pod でのモデルの実行をさらに詳しく調べることもできます。

より大きなモデルを使用する

ResNet-RS は、さまざまなサイズのモデル ファミリーを提供し、モデルが大きいほど、通常は計算コストが高くなります。詳細については、ResNet の再確認: トレーニングとスケーリング戦略の改善をご覧ください。

このコマンドの config_file を変更することで、トレーニングするモデルのサイズを選択できます。

(vm)$ python3 /usr/share/models/official/vision/beta/train.py \
  --experiment=resnet_rs_imagenet \
  --mode=train_and_eval \
  --model_dir=$MODEL_DIR \
  --tpu=$TPU_NAME \
  --config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs200_i256.yaml \
  --params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*"

利用可能な構成は VM 上の /usr/share/models/official/vision/beta/configs/experiments/ にあります。

Cloud TPU Pod でモデルのスケーリングする

Cloud TPU Pod でモデルをトレーニングするには、トレーニング スクリプトに変更を加える必要がある場合があります。詳細については、TPU Pod でのトレーニングをご覧ください。

Cloud TPU Pod を使用してモデルをスケーリングすると、より迅速に結果を得ることができます。完全にサポートされている ResNet-RS-50 モデルは、次の Pod スライスに対応しています。

  • v2-32
  • v3-32

Cloud TPU Pod では、トレーニングと評価が同時に行われます。

Cloud TPU Pod を使用したトレーニング

  1. 単一のデバイスでモデルをトレーニングするために作成した Cloud TPU リソースを削除します。

    (vm)$ gcloud compute tpus execution-groups delete resnet-rs-tutorial \
  --zone=europe-west4-a \
  --tpu-only

  2. accelerator-type パラメータを使用して、使用する Pod スライスを指定し、Pod 用の新しい Cloud TPU リソースを作成します。たとえば、次のコマンドは v2-32 Pod スライスを使用します。

    (vm)$ gcloud compute tpus execution-groups  create --name=resnet-rs-tutorial \
  --accelerator-type=v3-32  \
  --zone=europe-west4-a \
  --tf-version=2.12.0 \
  --tpu-only

    コマンドフラグの説明

    name
    作成する Cloud TPU の名前。
    accelerator-type
    作成する Cloud TPU のタイプ。
    zone
    Cloud TPU を作成するゾーン。
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
    tpu-only
    Cloud TPU のみを作成します。デフォルトでは、gcloud コマンドは VM と Cloud TPU を作成します。

  3. トレーニング スクリプトを実行します。

    (vm)$ python3 /usr/share/models/official/vision/beta/train.py \
--experiment=resnet_rs_imagenet \
--mode=train_and_eval \
--model_dir=$MODEL_DIR \
--tpu=$TPU_NAME \
--config_file=/usr/share/models/official/vision/beta/configs/experiments/image_classification/imagenet_resnetrs50_i160.yaml \
--params_override="task.train_data.input_path=$IMAGENET_DIR/train*, task.validation_data.input_path=$IMAGENET_DIR/valid*, trainer.train_steps=100"

    コマンドフラグの説明

    experiment
    実行するテストの名前。
    mode
    スクリプトを実行するモード。有効な値は、「train」、「eval」、または「train_and_eval」です。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    tpu
    使用する TPU の名前。
    config_file
    スクリプト構成ファイルへのパス。
    params_override
    スクリプト構成ファイルの設定をオーバーライドします。

これにより、ResNet-RS を 100 トレーニング ステップでトレーニングし、v3-8 TPU ノードでは 5 分以内に完了します。トレーニング スクリプトの出力には、次のようなテキストが含まれます。

{
  'train_loss': 1.435225,
  'train_accuracy': 0.00084427913
}

このトレーニング スクリプトも評価を実行します。評価出力には次のようなテキストが含まれます。

Run stats:
{
  'eval_loss': 0.861013,
  'eval_acc': 0.001,
  'train_loss': 1.435225,
  'train_acc': 0.00084427913,
  'step_timestamp_log': [
    'BatchTimestamp<batch_index: 0,
    timestamp: 1606330585.7613473>',
    'BatchTimestamp<batch_index: 500,
    timestamp: 1606330883.8486104>',
    'BatchTimestamp<batch_index: 1000,
    timestamp: 1606331119.515312>',
    'BatchTimestamp<batch_index: 1251,
    timestamp: 1606331240.7516596>'
  ],
  'train_finish_time': 1606331296.395158,
  'avg_exp_per_second': 1951.6983246161021
}

トレーニングと評価は一緒に行われます。合計 112,590 ステップ中、各エポックには、1,251 のトレーニング ステップと 48 の評価ステップがあります。

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud Shell で次のコマンドを使用して、Compute Engine VM と Cloud TPU を削除します。

    $ gcloud compute tpus execution-groups delete resnet-rs-tutorial \
  --zone=europe-west4-a

  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

    $ gcloud compute tpus execution-groups list \
 --zone=europe-west4-a

    次のような空の TPU のリストが表示されます。

       NAME             STATUS

  4. 次に示すように、gsutil を使用して Cloud Storage バケットを削除します。bucket-name を Cloud Storage バケットの名前に置き換えます。

    $ gsutil rm -r gs://bucket-name

次のステップ

TensorFlow Cloud TPU のチュートリアルでは通常、サンプル データセットを使用してモデルをトレーニングします。このトレーニングの結果は推論に使用できません。モデルを推論に使用するには、一般公開されているデータセットまたは独自のデータセットでデータをトレーニングします。Cloud TPU でトレーニングされた TensorFlow モデルは通常、データセットを TFRecord 形式にする必要があります。

データセット変換ツールのサンプルを使用して、画像分類データセットを TFRecord 形式に変換できます。画像分類モデルを使用しない場合は、自分でデータセットを TFRecord 形式に変換する必要があります。詳細については、TFRecord と tf.Example をご覧ください。

ハイパーパラメータ チューニング

データセットでモデルのパフォーマンスを向上させるには、モデルのハイパーパラメータを調整します。すべての TPU でサポートされているモデルに共通のハイパーパラメータに関する情報については、GitHub をご覧ください。モデルに固有のハイパーパラメータに関する情報については、各モデルのソースコードで確認できます。ハイパーパラメータ調整の詳細については、ハイパーパラメータ調整の概要とハイパーパラメータを調整するをご覧ください。

推論

モデルをトレーニングしたら、そのモデルを推論（予測）に使用できます。Cloud TPU 推論コンバータ ツールを使用して、Cloud TPU v5e での推論用の TensorFlow モデルを準備して最適化できます。Cloud TPU v5e での推論の詳細については、Cloud TPU v5e 推論の概要をご覧ください。