Vertex AI: 分散ハイパーパラメータ調整

1. 概要

このラボでは、Vertex AI でハイパーパラメータ調整と分散トレーニングを行う方法を学びます。このラボではモデルコードに TensorFlow を使用しますが、このコンセプトは別の ML フレームワークにも応用できます。

学習内容

次の方法を学習します。

• カスタム コンテナで分散トレーニングを使用してモデルをトレーニングする
• ハイパーパラメータ調整を自動化するためにトレーニング コードの複数のトライアルをリリースする

Google Cloud でこのラボを実行するための総費用は約 $6 USD です。

2. Vertex AI の概要

このラボでは、Google Cloud で利用できる最新の AI プロダクトを使用します。Vertex AI は Google Cloud 全体の ML サービスを統合してシームレスな開発エクスペリエンスを提供できるプロダクトです。これまで、AutoML でトレーニングしたモデルやカスタムモデルには、個別のサービスを介してアクセスする必要がありました。この新しいサービスは、それらを他の新しいプロダクトとともに 1 つの API へと結合します。既存のプロジェクトを Vertex AI に移行することもできます。ご意見やご質問がありましたら、サポートページからご連絡ください。

Vertex AI には、エンドツーエンドの ML ワークフローをサポートするさまざまなプロダクトが含まれています。このラボでは、主にトレーニングとワークベンチを使用します。

3. ユースケースの概要

このラボでは、ハイパーパラメータ調整を使用して、TensorFlow データセットにある馬と人のデータセットでトレーニングされた画像分類モデルに最適なパラメータを検出します。

ハイパーパラメータ調整

Vertex AI Training によるハイパーパラメータ調整では、選択したハイパーパラメータの値を使用して、指定された制限内でトレーニング アプリケーションのトライアルを複数回実行します。Vertex AI は、各トライアルの結果を追跡し、後続のトライアルのために調整を行います。

Vertex AI Training でハイパーパラメータ調整を使用するには、トレーニング コードに 2 つの変更を加える必要があります。

1. メインのトレーニング モジュールで、調整するハイパーパラメータごとにコマンドライン引数を定義します。
2. これらの引数に渡された値を使用して、アプリケーションのコード内の対応するハイパーパラメータを設定します。

分散トレーニング

GPU が 1 つの場合、TensorFlow はこのアクセラレータを使用して、追加作業なしでモデル トレーニングを高速化します。しかし、複数の GPU を使用することでさらなる効果を得たい場合は、複数のデバイスで計算を実行するための TensorFlow のモジュールである tf.distribute を使用する必要があります。

このラボでは、tf.distribute.MirroredStrategy を使用します。このコードは、わずかな変更でトレーニング アプリケーションに追加できます。この戦略では、マシン上の各 GPU にモデルのコピーを作成します。その後の勾配の更新は同期的に行われます。これは、各 GPU が、入力データの異なるスライスに対して、モデルを使用してフォワードパスとバックワード パスを計算することを意味しています。これらのスライスから計算された勾配は、すべての GPU に集約され、all-reduce と呼ばれるプロセスで平均化されます。モデル パラメータは、これらの平均勾配を使用して更新されます。

このラボを完了するために TensorFlow の分散トレーニングの仕組みを詳しく知る必要はありませんが、興味のある方は次の動画をご覧ください。

4.環境の設定

この Codelab を実行するには、課金が有効になっている Google Cloud Platform プロジェクトが必要です。プロジェクトを作成するには、こちらの手順を行ってください。

ステップ 1: Compute Engine API を有効にする

まだ有効になっていない場合は、[Compute Engine] に移動して [有効にする] を選択します。

ステップ 2: Container Registry API を有効にする

まだ有効になっていない場合は、[Container Registry] に移動して [有効にする] を選択します。これは、カスタム トレーニング ジョブのコンテナを作成する際に使用します。

ステップ 3: Vertex AI API を有効にする

Cloud Console の [Vertex AI] セクションに移動し、[Vertex AI API を有効にする] をクリックします。

ステップ 4: Vertex AI Workbench インスタンスを作成する

Cloud Console の [Vertex AI] セクションで [ワークベンチ] をクリックします。

Notebooks API をまだ有効にしていない場合は、有効にします。

有効にしたら、[マネージド ノートブック] をクリックします。

[新しいノートブック] を選択します。

ノートブックに名前を付けて、[詳細設定] をクリックします。

[詳細設定] で、アイドル状態でのシャットダウンを有効にして、シャットダウンまでの時間（分）を 60 に設定します。これにより、使用されていないノートブックが自動的にシャットダウンされるため、不要なコストが発生しません。

まだ有効になっていない場合は、[セキュリティ] で、[ターミナルを有効にする] を選択します。

詳細設定のその他の設定はそのままで構いません。

[作成] をクリックします。インスタンスがプロビジョニングされるまでに数分かかります。

インスタンスが作成されたら、[JUPYTERLAB を開く] を選択します。

新しいインスタンスを初めて使用するときに、認証が求められます。その場合は、UI で手順を行います。

5. トレーニング コードを作成する

まず、Launcher メニューから、ノートブック インスタンスでターミナル ウィンドウを開きます。

vertex-codelab という新しいディレクトリを作成し、そのディスクに移動します。

mkdir vertex-codelab
cd vertex-codelab

次のコマンドを実行して、トレーニング コードのディレクトリと、コードを追加する Python ファイルを作成します。

mkdir trainer
touch trainer/task.py

vertex-codelab ディレクトリに、次の行が作成されます。

+ trainer/
    + task.py

次に、作成した task.py ファイルを開き、以下のすべてのコードを貼り付けます。

import tensorflow as tf
import tensorflow_datasets as tfds
import argparse
import hypertune
import os

NUM_EPOCHS = 10
BATCH_SIZE = 64

def get_args():
  '''Parses args. Must include all hyperparameters you want to tune.'''

  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--learning_rate',
      required=True,
      type=float,
      help='learning rate')
  parser.add_argument(
      '--momentum',
      required=True,
      type=float,
      help='SGD momentum value')
  parser.add_argument(
      '--num_units',
      required=True,
      type=int,
      help='number of units in last hidden layer')
  args = parser.parse_args()
  return args

def preprocess_data(image, label):
  '''Resizes and scales images.'''

  image = tf.image.resize(image, (150,150))
  return tf.cast(image, tf.float32) / 255., label

def create_dataset(batch_size):
  '''Loads Horses Or Humans dataset and preprocesses data.'''

  data, info = tfds.load(name='horses_or_humans', as_supervised=True, with_info=True)

  # Create train dataset
  train_data = data['train'].map(preprocess_data)
  train_data  = train_data.shuffle(1000)
  train_data  = train_data.batch(batch_size)

  # Create validation dataset
  validation_data = data['test'].map(preprocess_data)
  validation_data  = validation_data.batch(batch_size)

  return train_data, validation_data

def create_model(num_units, learning_rate, momentum):
  '''Defines and compiles model.'''

  inputs = tf.keras.Input(shape=(150, 150, 3))
  x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu')(inputs)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(x)
  x = tf.keras.layers.MaxPooling2D((2, 2))(x)
  x = tf.keras.layers.Flatten()(x)
  x = tf.keras.layers.Dense(num_units, activation='relu')(x)
  outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
  model = tf.keras.Model(inputs, outputs)
  model.compile(
      loss='binary_crossentropy',
      optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=momentum),
      metrics=['accuracy'])
  return model

def main():
  args = get_args()

  # Create distribution strategy
  strategy = tf.distribute.MirroredStrategy()

  # Get data
  GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
  train_data, validation_data = create_dataset(GLOBAL_BATCH_SIZE)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model(args.num_units, args.learning_rate, args.momentum)

  # Train model
  history = model.fit(train_data, epochs=NUM_EPOCHS, validation_data=validation_data)

  # Define metric
  hp_metric = history.history['val_accuracy'][-1]

  hpt = hypertune.HyperTune()
  hpt.report_hyperparameter_tuning_metric(
      hyperparameter_metric_tag='accuracy',
      metric_value=hp_metric,
      global_step=NUM_EPOCHS)

if __name__ == "__main__":
    main()

コードをさらに調べて、分散トレーニングとハイパーパラメータ調整に固有のコンポーネントを見ていきましょう。

分散トレーニング

1. main() 関数で、MirroredStrategy オブジェクトが作成されます。次に、モデル変数の作成をストラテジーの範囲内に収めます。このステップでは、GPU 間でミラーリングする変数を TensorFlow に指定します。
2. バッチサイズは num_replicas_in_sync でスケールアップされます。バッチサイズのスケーリングは、TensorFlow で同期データ並列処理戦略を使用する場合のベスト プラクティスです。詳しくはこちらをご覧ください。

ハイパーパラメータ調整

1. このスクリプトは、hypertune ライブラリをインポートします。後でコンテナ イメージをビルドするときに、このライブラリをインストールする必要があります。
2. 関数 get_args() では、調整するハイパーパラメータごとにコマンドライン引数を定義します。この例で調整するハイパーパラメータは、学習率、オプティマイザーのモメンタム値、モデルの最終隠しレイヤのユニット数ですが、その他のパラメータの調整も自由に試してください。これらの引数に渡された値をコード内で使用し、対応するハイパーパラメータを設定します（たとえば、learning_rate = args.learning_rate に設定します）。
3. main() 関数の最後で、hypertune ライブラリを使用して、最適化する指標を定義します。TensorFlow では、Keras model.fit メソッドは History オブジェクトを返します。History.history 属性は、連続したエポックにおけるトレーニングの損失値と指標値の記録です。検証データを model.fit に渡した場合、History.history 属性には検証損失値と指標値も含まれます。たとえば、検証データを用いて 3 回のエポックでモデルをトレーニングし、指標として accuracy を提供した場合、History.history 属性は以下の辞書のようになります。

{
 "accuracy": [
   0.7795261740684509,
   0.9471358060836792,
   0.9870933294296265
 ],
 "loss": [
   0.6340447664260864,
   0.16712145507335663,
   0.04546636343002319
 ],
 "val_accuracy": [
   0.3795261740684509,
   0.4471358060836792,
   0.4870933294296265
 ],
 "val_loss": [
   2.044623374938965,
   4.100203514099121,
   3.0728273391723633
 ]

ハイパーパラメータ調整サービスでモデルの検証精度を最大化する値を見つけたい場合は、その指標を val_accuracy リストの最後のエントリ（または NUM_EPOCS - 1）として定義します。その後、この指標を HyperTune のインスタンスに渡します。hyperparameter_metric_tag には任意の文字列を選択できますが、ハイパーパラメータ調整ジョブを開始する際には、その文字列を再び使用する必要があります。

6. コードのコンテナ化

コードをコンテナ化する最初のステップは、Dockerfile の作成です。Dockerfile には、イメージの実行に必要なすべてのコマンドを含めます。必要なライブラリがすべてインストールされ、トレーニング コードのエントリ ポイントが設定されます。

ステップ 1: Dockerfile を作成する

ターミナルで、vertex-codelab ディレクトリに移動し、空の Dockerfile を作成します。

touch Dockerfile

vertex-codelab ディレクトリに、次の行が作成されます。

+ Dockerfile
+ trainer/
    + task.py

Dockerfile を開き、次の内容をコピーします。

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-7

WORKDIR /

# Installs hypertune library
RUN pip install cloudml-hypertune

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

この Dockerfile は Deep Learning Container TensorFlow Enterprise 2.7 GPU Docker イメージを使用します。Google Cloud の Deep Learning Containers には一般的な ML およびデータ サイエンス フレームワークが数多くプリインストールされています。この Dockerfile は、該当するイメージをダウンロードした後、トレーニング コードのエントリ ポイントを設定します。

ステップ 2: コンテナをビルドする

ターミナルで以下のように実行して、プロジェクトの環境変数を定義します。その際、your-cloud-project は実際のプロジェクト ID で置き換えてください。

ターミナルで gcloud config list --format 'value(core.project)' を実行すると、プロジェクト ID を取得できます。

PROJECT_ID='your-cloud-project'

Google Container Registry 内のコンテナ イメージの URI を示す変数を定義します。

IMAGE_URI="gcr.io/$PROJECT_ID/horse-human-codelab:latest"

Configure Docker

gcloud auth configure-docker

続いて、vertex-codelab ディレクトリのルートで次のように実行してコンテナをビルドします。

docker build ./ -t $IMAGE_URI

最後に、これを Google Container Registry に push します。

docker push $IMAGE_URI

ステップ 3: Cloud Storage バケットを作成する

トレーニング ジョブでは、ステージング バケットのパスを渡します。

ターミナルで次のコマンドを実行して、プロジェクトに新しいバケットを作成します。

BUCKET_NAME="gs://${PROJECT_ID}-hptune-bucket"
gsutil mb -l us-central1 $BUCKET_NAME

7. ハイパーパラメータ調整ジョブの開始

ステップ 1: ハイパーパラメータ調整を行うカスタム トレーニング ジョブを作成する

Launcher から新しい TensorFlow 2 Notebook を開きます。

Vertex AI Python SDK をインポートします。

from google.cloud import aiplatform
from google.cloud.aiplatform import hyperparameter_tuning as hpt

ハイパーパラメータ調整ジョブを開始するには、まずマシンタイプと Docker イメージを指定する worker_pool_specs を定義する必要があります。以下の仕様では、2 つの NVIDIA Tesla V100 GPU を搭載した 1 台のマシンを定義しています。

image_uri の {PROJECT_ID} は実際のプロジェクトに置き換えます。

# The spec of the worker pools including machine type and Docker image
# Be sure to replace PROJECT_ID in the "image_uri" with your project.

worker_pool_specs = [{
    "machine_spec": {
        "machine_type": "n1-standard-4",
        "accelerator_type": "NVIDIA_TESLA_V100",
        "accelerator_count": 2
    },
    "replica_count": 1,
    "container_spec": {
        "image_uri": "gcr.io/{PROJECT_ID}/horse-human-codelab:latest"
    }
}]

次に、parameter_spec を定義します。これは、最適化するパラメータを指定する辞書です。辞書のキーは、各ハイパーパラメータのコマンドライン引数に割り当てた文字列で、辞書の値はパラメータの仕様です。

各ハイパーパラメータにはタイプと調整サービスが試行する値の境界を定義する必要があります。ハイパーパラメータには、Double、Integer、Categorical、Discrete のタイプがあります。タイプに Double または Integer を選択した場合は、最小値と最大値を指定する必要があります。Categorical または Discrete を選択した場合は、値を入力する必要があります。Double タイプと Integer タイプでは、Scaling の値も必要になります。最適なスケールの選び方については、こちらの動画で詳しく紹介しています。

# Dictionary representing parameters to optimize.
# The dictionary key is the parameter_id, which is passed into your training
# job as a command line argument,
# And the dictionary value is the parameter specification of the metric.
parameter_spec = {
    "learning_rate": hpt.DoubleParameterSpec(min=0.001, max=1, scale="log"),
    "momentum": hpt.DoubleParameterSpec(min=0, max=1, scale="linear"),
    "num_units": hpt.DiscreteParameterSpec(values=[64, 128, 512], scale=None)
}

定義する最後の仕様は metric_spec です。これは最適化する指標を表す辞書です。辞書のキーはトレーニング アプリケーションのコードで設定した hyperparameter_metric_tag で、値は最適化の目標です。

# Dicionary representing metrics to optimize.
# The dictionary key is the metric_id, which is reported by your training job,
# And the dictionary value is the optimization goal of the metric.
metric_spec={'accuracy':'maximize'}

spec を定義してから、CustomJob を作成します。これは、ハイパーパラメータ調整の各トライアルでジョブを実行するために使用される共通の spec です。

{YOUR_BUCKET} は、前に作成したバケットに置き換える必要があります。

# Replace YOUR_BUCKET
my_custom_job = aiplatform.CustomJob(display_name='horses-humans',
                              worker_pool_specs=worker_pool_specs,
                              staging_bucket='gs://{YOUR_BUCKET}')

次に、HyperparameterTuningJob を作成して実行します。

hp_job = aiplatform.HyperparameterTuningJob(
    display_name='horses-humans',
    custom_job=my_custom_job,
    metric_spec=metric_spec,
    parameter_spec=parameter_spec,
    max_trial_count=6,
    parallel_trial_count=2,
    search_algorithm=None)

hp_job.run()

注意すべき引数がいくつかあります。

• max_trial_count: サービスが実行するトライアルの数に上限を設ける必要があります。トライアルの回数を増やすと一般的に良い結果が得られますが、収穫逓減のポイントがあり、それ以降はトライアルの回数を増やしても最適化しようとしている指標にほとんど影響がなくなります。少ないトライアル回数から始めて、選択したハイパーパラメータの影響力を把握してから、スケールアップするのがベスト プラクティスです。
• parallel_trial_count: パラレル トライアルをご利用の場合、サービスは複数のトレーニング処理クラスタをプロビジョニングします。パラレル トライアルの回数を増やすと、ハイパーパラメータ調整ジョブの実行時間が短縮されますが、ジョブ全体の効果が低下する可能性がありますこれは、デフォルトの調整戦略では、過去のトライアルの結果を後続のトライアルでの値の割り当てに使用するためです。
• search_algorithm: 検索アルゴリズムは、グリッド、ランダム、デフォルト（なし）のいずれかに設定できます。デフォルトのオプションは、ベイズ最適化を適用して有効なハイパーパラメータ値の領域を検索するもので、推奨されるアルゴリズムとなっています。詳しくは、こちらをご覧ください。

ジョブが開始されると、UI の [ハイパーパラメータ調整ジョブ] タブでステータスを確認できるようになります。

ジョブが完了したら、トライアルの結果を表示して並べ替えて、ハイパーパラメータ値の最適な組み合わせを見つけます。

🎉 お疲れさまでした 🎉

Vertex AI を使って次のことを行う方法を学びました。

• 分散トレーニングでハイパーパラメータ調整ジョブを実行する

Vertex AI のさまざまな機能の詳細については、こちらのドキュメントをご覧ください。

8. クリーンアップ

ノートブックは、アイドル状態で 60 分が経過するとタイムアウトするように構成されています。このため、インスタンスのシャットダウンを心配する必要はありません。インスタンスを手動でシャットダウンする場合は、Console で [Vertex AI] の [ワークベンチ] セクションにある [停止] ボタンをクリックします。ノートブックを完全に削除する場合は、[削除] ボタンをクリックします。

ストレージ バケットを削除するには、Cloud Console のナビゲーション メニューで [ストレージ] を移動してバケットを選択し、[削除] をクリックします。