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
  • クラウド ストレージ

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

始める前に

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

  1. Google Cloud アカウントにログインします。Google Cloud を初めて使用する場合は、アカウントを作成して、実際のシナリオでの Google プロダクトのパフォーマンスを評価してください。新規のお客様には、ワークロードの実行、テスト、デプロイができる無料クレジット $300 分を差し上げます。

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

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

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

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

リソースを設定する

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

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

    Cloud Shell を開く

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

    export PROJECT_ID=project-id

  3. Cloud TPU を作成するプロジェクトを使用するように gcloud コマンドライン ツールを構成します。

    gcloud config set project ${PROJECT_ID}

    このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した GCP 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 -b on gs://bucket-name

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

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

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

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

    コマンドフラグの説明

    vm-only
    VM のみを作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    disk-size
    gcloud compute tpus execution-groups コマンドで作成された VM のハードディスクのサイズ（GB）。
    machine-type
    作成する Compute Engine VM のマシンタイプ。
    tf-version
    gcloud compute tpusexecution-groups で VM にインストールされる TensorFlow のバージョン。

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

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

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

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

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

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
(vm)$ export PYTHONPATH="$PYTHONPATH:/usr/share/models/"

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

fake_imagenet を使用して ResNet モデルのトレーニングと評価を行う

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 リソースを起動します。

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

    コマンドフラグの説明

    tpu-only
    VM を作成せずに Cloud TPU を作成します。デフォルトでは、gcloud compute tpus execution-groups コマンドは VM と Cloud TPU を作成します。
    accelerator-type
    作成する Cloud TPU のタイプ。
    name
    作成する Cloud TPU の名前。
    zone
    Cloud TPU を作成するゾーン。
    tf-version
    VM にインストールする TensorFlow のバージョン。

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

    (vm)$ export TPU_NAME=resnet-tutorial

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

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

  4. ResNet-50 モデル ディレクトリに移動します。

    (vm)$ cd /usr/share/models/official/vision/image_classification/resnet/

  5. トレーニング スクリプトを実行します。これは、fake_imagenet データセットを使用し、ResNet を 1 エポック、トレーニングします。

    (vm)$ python3 resnet_ctl_imagenet_main.py \
 --tpu=${TPU_NAME} \
 --model_dir=${MODEL_DIR} \
 --data_dir=${DATA_DIR} \
 --batch_size=1024 \
 --steps_per_loop=500 \
 --train_epochs=1 \
 --use_synthetic_data=false \
 --dtype=fp32 \
 --enable_eager=true \
 --enable_tensorboard=true \
 --distribution_strategy=tpu \
 --log_steps=50 \
 --single_l2_loss_op=true \
 --use_tf_function=true

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    data_dir
    トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
    batch_size
    トレーニングのバッチサイズ。
    steps_per_loop
    CPU に状態を保存する前に実行するトレーニング ステップの数。トレーニング ステップは、サンプルの 1 つのバッチの処理です。これには、フォワードパスと逆伝播の両方が含まれます。
    train_epochs
    モデルをトレーニングするためのエポックの数。
    use_synthetic_data
    トレーニングで合成データを使用するには、true に設定します。
    dtype
    トレーニングに使用するデータの種類。
    enable_eager
    TensorFlow 積極的実行（Eager Execution）を有効にします。
    enable_tensorboard
    TensorBoard を有効にします。
    distribution_strategy
    Cloud TPU で ResNet モデルをトレーニングするには、tpu に設定します。
    log_steps
    examples per second などのタイミング情報をログに記録する前に実行するトレーニング ステップの数。
    single_l2_loss_op
    Keras のレイヤごとの L2 損失を使用する代わりに、true に設定して連結された重みで L2_loss を計算します。
    use_tf_function
    トレーニングとテストのステップを tf.function 内にラップするには、true に設定します。

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

{
  '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 をトレーニングするには、次のスクリプトに示すように 90 エポックで実行します。トレーニングと評価は一緒に行われます。合計 112,590 ステップ中、各エポックには、1,251 のトレーニング ステップと 48 の評価ステップがあります。

(vm)$ python3 resnet_ctl_imagenet_main.py \
    --tpu=${TPU_NAME} \
    --model_dir=${MODEL_DIR} \
    --data_dir=${DATA_DIR} \
    --batch_size=1024 \
    --steps_per_loop=500 \
    --train_epochs=90 \
    --use_synthetic_data=false \
    --dtype=fp32 \
    --enable_eager=true \
    --enable_tensorboard=true \
    --distribution_strategy=tpu \
    --log_steps=50 \
    --single_l2_loss_op=true \
    --use_tf_function=true

コマンドフラグの説明

tpu
Cloud TPU の名前。
model_dir
トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
data_dir
トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
batch_size
トレーニングのバッチサイズ。
steps_per_loop
CPU に状態を保存する前に実行するトレーニング ステップの数。トレーニング ステップは、サンプルの 1 つのバッチの処理です。これには、フォワードパスと逆伝播の両方が含まれます。
train_epochs
モデルをトレーニングするためのエポックの数。
use_synthetic_data
トレーニングに合成データを使用するかどうか。
dtype
トレーニングに使用するデータの種類。
enable_eager
TensorFlow 積極的実行（Eager Execution）を有効にします。
enable_tensorboard
TensorBoard を有効にします。
distribution_strategy
Cloud TPU で ResNet モデルをトレーニングするには、tpu に設定します。
log_steps
examples per second などのタイミング情報をログに記録する前に実行するトレーニング ステップの数。
single_l2_loss_op
Keras のレイヤごとの L2 損失を使用する代わりに、true に設定して連結された重みで L2_loss を計算します。
use_tf_function
トレーニングとテストのステップを tf.function 内にラップするには、true に設定します。

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

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

Cloud TPU Pod を使用してモデルをスケーリングする

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

  • v2-32
  • v3-32

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

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

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

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

  2. 使用する Pod スライスを指定するための accelerator-typeパラメータを使用して、gcloud compute tpus execution-groups コマンドを実行します。たとえば、次のコマンドは v3-32 の Pod スライスを使用します。

    (vm)$ gcloud compute tpus execution-groups  create --name=resnet-tutorial \
  --accelerator-type=v3-32  \
  --zone=europe-west4-a \
  --tf-version=2.5.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)$ export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x-pod

  4. モデルをトレーニングします。

    (vm)$ python3 resnet_ctl_imagenet_main.py \
  --tpu=${TPU_NAME} \
  --model_dir=${MODEL_DIR} \
  --data_dir=${DATA_DIR} \
  --batch_size=4096 \
  --steps_per_loop=500 \
  --train_epochs=1 \
  --use_synthetic_data=false \
  --dtype=fp32 \
  --enable_eager=true \
  --enable_tensorboard=true \
  --distribution_strategy=tpu \
  --log_steps=50 \
  --single_l2_loss_op=true \
  --use_tf_function=true

    コマンドフラグの説明

    tpu
    Cloud TPU の名前。
    model_dir
    トレーニング中にチェックポイントとサマリーが保存される Cloud Storage バケット。既存のフォルダを使用して、同じサイズの TPU と TensorFlow バージョンで以前に生成されたチェックポイントを読み込むことができます。
    data_dir
    トレーニング入力の Cloud Storage パス。この例では、fake_imagenet データセットに設定されています。
    batch_size
    トレーニングのバッチサイズ。
    steps_per_loop
    CPU に状態を保存する前に実行するトレーニング ステップの数。トレーニング ステップは、サンプルの 1 つのバッチの処理です。これには、フォワード パスと逆伝播の両方が含まれます。
    train_epochs
    モデルをトレーニングするためのエポックの数。
    use_synthetic_data
    トレーニングで合成データを使用するには、true に設定します。
    dtype
    トレーニングに使用するデータの種類。
    enable_eager
    TensorFlow の積極的実行（Eager Execution）を有効にするには、true に設定します。
    enable_tensorboard
    TensorBoard を有効にするには、true に設定します。
    distribution_strategy
    Cloud TPU で ResNet モデルをトレーニングするには、tpu に設定します。
    log_steps
    examples per second などのタイミング情報をログに記録する前に実行するトレーニング ステップの数。
    single_l2_loss_op
    レイヤごとの L2 損失を使用する代わりに、連結された重みで L2_loss を計算します。
    use_tf_function
    トレーニングとテストのステップを tf.function 内にラップします。

この手順は、fake_imagnet データセットに対してモデルを 1 エポックまでトレーニング（合計 312 トレーニング ステップ、12 評価ステップ）します。このトレーニングは、v3-32 Cloud TPU 上で約 2 分かかります。トレーニング スクリプトによって次のようなテキストが表示されます。

step: 312
steps_per_second: 1.38
{
  'train_loss': 0.3755003,
  'train_accuracy': 0.000983605
}

また、このスクリプトによって評価が行われ、次のようなテキストが表示されます。

step: 312
evaluation metric: {
  'test_loss': 0.20280758,
  'test_accuracy': 0.001
}

クリーンアップ

このチュートリアルで使用したリソースについて、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-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

次のステップ

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

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

ハイパーパラメータ調整

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

推論

モデルをトレーニングしたら、そのモデルを推論（予測）に使用できます。AI Platform は、機械学習モデルを開発、トレーニング、デプロイするためのクラウドベースのソリューションです。モデルをデプロイすれば、AI Platform Prediction サービスを使用できるようになります。

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