Cloud TPU(TF 2.x)での ResNet のトレーニング


このチュートリアルでは、tf.distribute.TPUStrategy を使用して Cloud TPU で Keras ResNet モデルをトレーニングする方法を説明します。

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. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

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

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 に認証情報を使用した API の呼び出しを許可します。

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

    サービス アカウントにより、Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。

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

  6. データセットは、独自に用意するか、fake_imagenet を使用します。

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

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

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

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

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

    完全な ImageNet データセットを使用する場合は、ImageNet データセットのダウンロード、前処理、アップロードをご覧ください。

  7. gcloud コマンドを使用して TPU リソースを起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm create resnet-tutorial \
    --zone=us-central2-b \
    --accelerator-type=v4-8 \
    --version=tpu-vm-tf-2.16.1-pjrt
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン

    TPU ノード

    gcloud compute tpus execution-groups create \
    --project=${PROJECT_ID} \
    --zoneus-central2-b \
    --name=resnet-tutorial \
    --disk-size=300 \
    --machine-type=n1-standard-16 \
    --accelerator-type=v3-8 \
    --tf-version=2.12.0

    コマンドフラグの説明

    project
    : Google Cloud プロジェクト ID
    zone
    Cloud TPU を作成するゾーン
    name
    作成する Cloud TPU の名前。
    disk-size
    gcloud コマンドで作成された VM のハードディスクのサイズ(GB)。
    machine-type
    作成する Compute Engine VM のマシンタイプ
    accelerator-type
    作成する Cloud TPU のタイプ
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。

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

  8. 自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェル プロンプトが username@projectname から username@vm-name に変わります。

    TPU VM

    $ gcloud compute tpus tpu-vm ssh resnet-tutorial --zone=us-central2-b
    

    TPU ノード

    $ gcloud compute ssh resnet-tutorial --zone=us-central2-b
    

    これらの手順を続行する場合は、VM セッション ウィンドウで、(vm)$ で始まる各コマンドを実行します。

  9. Cloud TPU 名の変数を設定します。

    TPU VM

    (vm)$ export TPU_NAME=local
    

    TPU ノード

    (vm)$ export TPU_NAME=resnet-tutorial
    
  10. Cloud Storage バケットの変数を設定する

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

    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
    (vm)$ export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    
    

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

  11. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
    
  12. TensorFlow の要件をインストールします。

    使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。

    TPU VM

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt 

    TPU ノード

    (vm)$ pip3 install --user -r /usr/share/models/official/requirements.txt
    
  13. ResNet トレーニング スクリプトには、追加のパッケージが必要です。この時点でインストールしておきます。

    (vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3
    
  14. モデルが保存されているディレクトリに移動します。

    TPU VM

    (vm)$ cd /usr/share/tpu/tensorflow/resnet50_keras
    

    TPU ノード

    (vm)$ cd /usr/share/models
    
  15. PYTHONPATH 環境変数を設定します。

    TPU VM

    (vm)$ /usr/share/tpu/tensorflow/resnet50_keras"
    

    TPU ノード

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    
  16. トレーニング スクリプトを実行します。これは fake_imagenet データセットを使用し、ResNet を 100 ステップでトレーニングします。

    TPU VM

    (vm)$ resnet50.py --tpu=local --data=gs://cloud-tpu-test-datasets/fake_imagenet

    Command flag descriptions

    tpu
    The name of your TPU.
    data
    Specifies the directory where checkpoints and summaries are stored during model training. If the folder is missing, the program creates one. When using a Cloud TPU, the model_dir must be a Cloud Storage path (gs://...). You can reuse an existing folder to load current checkpoint data and to store additional checkpoints as long as the previous checkpoints were created using TPU of the same size and TensorFlow version.

    TPU ノード

    (vm)$ python3 official/vision/train.py \
    --tpu=${TPU_NAME} \
    --experiment=resnet_imagenet \
    --mode=train_and_eval \
    --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
    --model_dir=${MODEL_DIR} \
    --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
    

    コマンドフラグの説明

    tpu
    TPU の名前。
    model_dir
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリを指定します。該当するフォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir は Cloud Storage のパスである必要があります(gs://...)。以前に作成されたチェックポイントの TPU のサイズと TensorFlow のバージョンが同じであれば、既存のフォルダを再利用して現在のチェックポイント データを読み込み、追加のチェックポイントを保存できます。

これにより ResNet が 100 ステップでトレーニングされ、v3-8 TPU ノードでは約 3 分で完了します。100 ステップの最後に、次のような出力が表示されます。

I0624 17:04:26.974905 140457742666816 controller.py:290]  eval | step:    100 | eval time:   23.3 sec | output:
    {'accuracy': 0.0010141226,
     'top_5_accuracy': 0.0051457332,
     'validation_loss': 8.448798}
 eval | step:    100 | eval time:   23.3 sec | output:
    {'accuracy': 0.0010141226,
     'top_5_accuracy': 0.0051457332,
     'validation_loss': 8.448798}

これで単一デバイスのトレーニング サンプルは完了しました。現在の単一デバイスの TPU リソースを削除するには、次の手順を使用します。

  1. Compute Engine インスタンスから接続を切断します。

    (vm)$ exit
    

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

  2. TPU リソースを削除します。

    TPU VM

    $ gcloud compute tpus tpu-vm delete resnet-tutorial \
    --zone=us-central2-b
    

    コマンドフラグの説明

    zone
    Cloud TPU が存在するゾーン

    TPU ノード

    $ gcloud compute tpus execution-groups delete resnet-tutorial \
    --zone=us-central2-b
    

    コマンドフラグの説明

    zone
    削除する TPU を含むゾーン

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

Cloud TPU Pod を使用したモデルのスケーリング

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

TPU Pod のトレーニング

このセクションでは、Pod のトレーニング用の Cloud Storage バケットと 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 us-central2 gs://bucket-name
    

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

    バケットのロケーションは、TPU VM と同じリージョンとなる必要があります。

  6. データセットは、独自に用意するか、fake_imagenet を使用します。

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

    デフォルトの Pod トレーニングでは、ImageNet の完全版データセットの fnet_imagenet と呼ばれるデモバージョンにアクセスします。このデモバージョンを使用すると、ImageNet データベース全体に対してモデルをトレーニングするときに通常伴うストレージと時間の要件を削減すると同時に Pod のトレーニングをテストできます。

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

    完全な ImageNet データセットを使用する場合は、ImageNet データセットのダウンロード、前処理、アップロードをご覧ください。

  7. gcloud コマンドを使用して Cloud TPU リソースを起動します。

    使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。このチュートリアルでは、v3-32 Pod を指定します。他の Pod オプションについては、TPU のバージョンをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm create resnet-tutorial \
    --zone=us-central2-b \
    --accelerator-type=v4-32 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン

    TPU ノード

    $ gcloud compute tpus execution-groups create \
    --zone=us-central2-b \
    --name=resnet-tutorial \
    --accelerator-type=v3-32 \
    --tf-version=2.12.0
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    name
    作成する Cloud TPU の名前。
    accelerator-type
    作成する Cloud TPU のタイプ
    tf-version
    Tensorflow gcloud のバージョンが VM にインストールされます。
  8. 自動的に Compute Engine インスタンスにログインしない場合は、次の ssh コマンドを実行してログインします。VM にログインすると、シェル プロンプトが username@projectname から username@vm-instance-name に変わります。

    TPU VM

    $ gcloud compute tpus tpu-vm ssh resnet-tutorial --zone=us-central2-b
    

    TPU ノード

    $ gcloud compute ssh resnet-tutorial --zone=us-central2-b
    

    これらの手順を続行する場合は、VM セッション ウィンドウで、(vm)$ で始まる各コマンドを実行します。

  9. Cloud TPU の設定変数をエクスポートします。

    (vm)$ export ZONE=us-central2-b
    (vm)$ export STORAGE_BUCKET=gs://bucket-name
    
    (vm)$ export TPU_NAME=resnet-tutorial
    (vm)$ export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    (vm)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x-pod
    

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

  10. ResNet トレーニング スクリプトには、追加のパッケージが必要です。この時点でインストールしておきます。

    TPU VM

    (vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3 

    TPU ノード

    (vm)$ pip3 install --user tensorflow-model-optimization>=0.1.3 
  11. TensorFlow の要件をインストールします。

    使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。

    TPU VM

    (vm)$ pip3 install -r /usr/share/tpu/models/official/requirements.txt
    

    TPU ノード

    (vm)$ pip3 install --user -r /usr/share/models/official/requirements.txt
    
  12. PYTHONPATH 環境変数を設定します。

    TPU VM

    (vm)$ export PYTHONPATH="PYTHONPATH=/usr/share/tpu/tensorflow/resnet50_keras"
    (vm)$ export TPU_LOAD_LIBRARY=0
    

    TPU ノード

    (vm)$ export PYTHONPATH="${PYTHONPATH}:/usr/share/models"
    
  13. モデルが保存されているディレクトリに移動します。

    TPU VM

    (vm)$ cd /usr/share/tpu/tensorflow/resnet50_keras
    

    TPU ノード

    (vm)$ cd /usr/share/models
    
  14. モデルをトレーニングします。

    (vm)$ resnet50.py --tpu=${TPU_NAME} --data=gs://cloud-tpu-test-datasets/fake_imagenet
     

    コマンドフラグの説明

    tpu
    TPU の名前。
    data
    モデルのトレーニング中にチェックポイントとサマリーが保存されるディレクトリを指定します。該当するフォルダがない場合は、プログラムによって作成されます。Cloud TPU を使用する場合、model_dir を Cloud Storage パスにする必要があります(gs://...)。以前のチェックポイントが、同じサイズの Cloud TPU と TensorFlow バージョンを使用して作成されていれば、既存のフォルダを再利用して現在のチェックポイント データを読み込んで追加のチェックポイントを保存できます。

この手順では、fake_imagenet データセットで 100 トレーニング ステップ、13 評価ステップでモデルをトレーニングします。このトレーニングは、v3-32 Cloud TPU 上で約 2 分かかります。トレーニングと評価が完了すると、次のようなメッセージが表示されます。

{'accuracy': 0.0009716797,
     'learning_rate': 0.10256411,
     'top_5_accuracy': 0.0049560545,
     'training_loss': 8.5587225}
train | step:    100 | steps/sec:    1.2 | output:
    {'accuracy': 0.0009716797,
     'learning_rate': 0.10256411,
     'top_5_accuracy': 0.0049560545,
     'training_loss': 8.5587225}

eval | step:    100 | eval time:   24.8 sec | output:
    {'accuracy': 0.0010141226,
     'top_5_accuracy': 0.004356971,
     'validation_loss': 8.50038}
 eval | step:    100 | eval time:   24.8 sec | output:
    {'accuracy': 0.0010141226,
     'top_5_accuracy': 0.004356971,
     'validation_loss': 8.50038}

クリーンアップ

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

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

    (vm)$ exit
    

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

  2. Cloud TPU と Compute Engine リソースを削除します。リソースの削除に使用するコマンドは、TPU VM または TPU ノードのどちらを使用するかによって異なります。詳細については、システム アーキテクチャをご覧ください。

    TPU VM

    $ gcloud compute tpus tpu-vm delete resnet-tutorial \
    --zone=us-central2-b
    

    TPU ノード

    $ gcloud compute tpus execution-groups delete resnet-tutorial \
    --zone=us-central2-b
    
  3. gcloud compute tpus execution-groups list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。次のコマンドの出力には、このチュートリアルで作成したリソースを含めないでください。

    TPU VM

    $ gcloud compute tpus tpu-vm list --zone=us-central2-b

    TPU ノード

    $ gcloud compute tpus execution-groups list --zone=us-central2-b
  4. 次に示すように gsutil を実行します。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 推論の概要をご覧ください。

  • データセット変換チュートリアルに従って、fake_imagenet データセットまたは ImageNet データセットの代わりに独自のデータを使用してトレーニングと評価を行う方法を学習します。 このチュートリアルでは、画像分類データ変換プログラムのサンプル スクリプトを使用して、画像分類用の生のデータセットを Cloud TPU Tensorflow モデルで使用される TFRecord に変換する方法について説明します。
  • 独自の画像データを使用して画像分類モデルを実行する方法を示す Cloud TPU colab を実行します。
  • 他の Cloud TPU チュートリアルを確認する。
  • TensorBoard の TPU モニタリング ツールの使用方法を学習する。
  • Cloud TPU と GKE を使用して ResNet をトレーニングする方法を確認します。