Vertex AI 上の TensorFlow による分散型トレーニングとハイパーパラメータ調整
Google Cloud Japan Team
※この投稿は米国時間 2021 年 8 月 5 日に、Google Cloud blog に投稿されたものの抄訳です。
モデルのハイパーパラメータにどのような値を選択するかで、大きな違いが生まれます。少数のハイパーパラメータを調整するだけであれば、手動で実験を行うことができるかもしれません。しかし、ディープ ラーニング モデルの場合、アーキテクチャやオプティマイザーのハイパーパラメータを頻繁に変更したり、最適なバッチサイズや学習率を見つけようとしたりすることが多く、すぐにこうした実験を大規模に自動化する必要が生じます。
この記事では、Vertex Training でハイパーパラメータ調整ジョブを実行し、ML モデルの最適なハイパーパラメータ値を発見する方法の例を説明します。また、トレーニング プロセスを高速化するために、tf.distribute Python モジュールを利用して、1 台のマシン上の複数の GPU にコードを分散させます。このチュートリアルのコードはすべて、このノートブックに記載されています。
ハイパーパラメータ調整サービスを利用するには、トレーニング アプリケーション コードとカスタム トレーニング ジョブ リクエストで、調整するハイパーパラメータを定義する必要があります。トレーニング アプリケーションのコードでは、各ハイパーパラメータのコマンドライン引数を定義し、それらの引数で渡された値を使って、コード内で対応するハイパーパラメータを設定します。また、cloudml-hypertune Python パッケージを使って、最適化したい指標を Vertex AI に報告する必要があります。
この例では TensorFlow を使用していますが、PyTorch や XGBoost など、お好みのフレームワークで書かれたモデルで Vertex Training を使用できます。
tf.distribute.Strategy API の使用
GPU が 1 つの場合、TensorFlow はこのアクセラレータを使用して、追加作業なしでモデル トレーニングを高速化します。しかし、1 台のマシンや複数のマシン（それぞれが複数の GPU を持つ可能性がある）で複数の GPU を使用することでさらなる効果を得たい場合は、複数のデバイスで計算を実行するための TensorFlow のモジュールである tf.distribute を使用する必要があります。分散トレーニングを始めるための最もシンプルな方法は、1 台のマシンに複数の GPU デバイスを搭載することです。tf.distribute.Strategy API からの TensorFlow 配信ストラテジーは、すべての GPU にわたるデータ配信の調整と勾配の更新を管理します。
tf.distribute.MirroredStrategy は、わずかなコード変更で使用できる同期型のデータ並列処理ストラテジーです。この戦略では、マシン上の各 GPU にモデルのコピーを作成します。その後の勾配の更新は同期的に行われます。これは、各ワーカー デバイスが、入力データの異なるスライスに対して、モデルを使用してフォワードパスとバックワード パスを計算することを意味しています。これらのスライスから計算された勾配は、すべての GPU に集約され、all-reduce と呼ばれるプロセスで低減されます。オプティマイザーは、これらの低減された勾配でパラメータの更新を行うことで、デバイスの同期を維持します。
tf.distribute.Strategy API を使用するための最初のステップは、ストラテジー オブジェクトの作成です。
次に、モデル変数の作成をストラテジーの範囲内に収める必要があります。このステップは、GPU デバイス間でミラーリングする変数を MirroredStrategy に伝えるため、非常に重要です。
最後に、バッチサイズを GPU の数でスケーリングします。tf.distribute.Strategy API と tf.data を使って分散型トレーニングを行う場合、バッチサイズはグローバル バッチサイズを参照するようになりました。つまり、バッチサイズを 16 に設定し、GPU を 2 つ持っている場合、各マシンは 1 ステップにつき 8 例を処理します。ここでは、16 がグローバル バッチサイズ、8 がレプリカごとのバッチサイズと呼ばれています。GPU を最大限に活用するためには、バッチサイズをレプリカの数でスケーリングする必要があります。
なお、コードの配信は任意です。なお、複数の GPU を使用しない場合でも、次項の手順に沿ってハイパーパラメータ調整サービスを利用できます。
ハイパーパラメータ調整のためのトレーニング コードの更新
Vertex Training でハイパーパラメータ調整を使用するには、トレーニング コードに 2 つの変更を加える必要があります。
まず、メインのトレーニング モジュールで、調整したいハイパーパラメータごとにコマンドライン引数を定義します。次に、これらの引数で渡された値を使って、アプリケーションのコードで対応するハイパーパラメータを設定します。
例えば、学習率、オプティマイザーのモメンタム値、モデルの最終隠しレイヤのニューロン数を調整したいとします。以下の関数に示すように、argparse を使ってコマンドラインの引数を解析できます。
これらの引数には好きな名前を付けることができますが、アプリケーションのコードで対応するハイパーパラメータを設定するためには、これらの引数で渡された値を使用する必要があります。例えば、オプティマイザーは次のようになります。
調整したいハイパーパラメータがわかったところで、最適化する指標を決める必要があります。ハイパーパラメータ調整サービスが複数回のトライアルを行った後、モデルに選択されるハイパーパラメータの値は、選択された指標を最大化（または最小化）する組み合わせになります。
この指標を報告するには、cloudml-hypertune ライブラリを使用します。このライブラリは、どのフレームワークでも使用できます。
TensorFlow では、keras model.fit メソッドは History オブジェクトを返します。
History.history 属性は、連続したエポックにおけるトレーニングの損失値と指標値の記録です。検証データを model.fit に渡した場合、History.history 属性には検証損失値と指標値も含まれます。
たとえば、検証データを用いて 3 回のエポックでモデルをトレーニングし、指標としてaccuracy を提供した場合、History.history 属性は以下の辞書のようになります。
検証精度を最大化する学習率、動向、ユニット数の値を選択するために、「val_accuracy」リストの最後のエントリ（または NUM_EPOCS - 1）としてに指標を定義します。
次に、この指標を HyperTune のインスタンスに渡し、HyperTune は各トレーニングの終了時にその値を Vertex AI に報告します。
これで完了です。以上の 2 つの簡単なステップで、トレーニング アプリケーションが完成します。
ハイパーパラメータ調整ジョブの開始
トレーニングのアプリケーション コードを修正したら、ハイパーパラメータの調整ジョブを開始します。この例では、Python SDK を使ってジョブを起動する方法を示していますが、Cloud Console UI を使用することもできます。
トレーニングのアプリケーション コードが、カスタム コンテナとしてパッケージ化されていることを確認する必要があります。その方法がよくわからない場合は、詳細な手順についてこのチュートリアルを参照してください。
ノートブックで、ランチャーから新しい Python 3 ノートブックを作成します。
ノートブックで、以下をセル内で実行し、Vertex AI SDK をインストールします。セルが終了したら、カーネルを再起動します。
カーネルを再起動した後、SDK をインポートします。
ハイパーパラメータの調整ジョブを起動するには、まずマシンタイプと Docker イメージを指定する worker_pool_specs を定義する必要がありま以下の仕様では、2 つの NVIDIA T4 Tensor Core GPU を搭載した 1 台のマシンを定義しています。
次に、最適化したいパラメータを指定した辞書である parameter_spec を定義します。辞書のキーは、各ハイパーパラメータのコマンドライン引数に割り当てた文字列で、辞書の値はパラメータの仕様です。
各ハイパーパラメータにはタイプと調整サービスが試行する値の境界を定義する必要があります。タイプに Double や Integer を選択した場合は、最小値と最大値を指定する必要があります。また、Categorical または Discrete を選択した場合は、値を入力する必要があります。Double タイプと Integer タイプでは、Scaling の値も必要になります。最適なスケールの選び方については、こちらの動画で詳しく紹介しています。最後に定義する spec は、最適化する指標を表す辞書である metric_spec です。辞書のキーはトレーニング アプリケーションのコードで設定した hyperparameter_metric_tag で、値は最適化の目標です。
spec を定義してから、CustomJob を作成します。これは、ハイパーパラメータ調整の各トライアルでジョブを実行するために使用される共通の spec です。
{YOUR_BUCKET} を、ステージング用のプロジェクトのバケットに置き換える必要があります。
最後に、HyperparameterTuningJob を作成し、実行します。
注意すべき引数がいくつかあります。
max_trial_count: サービスが実行するトライアルの数に上限を設ける必要があります。トライアルの回数を増やすと一般的に良い結果が得られますが、収穫逓減のポイントがあり、それ以降はトライアルの回数を増やしても最適化しようとしている指標にほとんど影響がなくなります。少ないトライアル回数から始めて、選択したハイパーパラメータの影響力を把握してから、スケールアップするのがベスト プラクティスです。
parallel_trial_count: パラレル トライアルをご利用の場合、サービスは複数のトレーニング処理クラスタをプロビジョニングします。ジョブ作成時に指定するワーカープールの spec は、個々のトレーニング クラスタに使用されます。 パラレル トライアルの回数を増やすと、ハイパーパラメータ調整ジョブの実行時間が短縮されますが、ジョブ全体の効果が低下する可能性がありますこれは、デフォルトの調整戦略では、過去のトライアルの結果を後続のトライアルでの値の割り当てに使用するためです。
search_algorithm: 検索アルゴリズムは、グリッド、ランダム、デフォルト（なし）のいずれかに設定できます。グリッド検索はハイパーパラメータを徹底的に検索しますが、高次元空間では実現不可能です。ランダム検索では、検索空間をランダムにサンプリングします。ランダム検索の欠点は、過去のテストの情報を次の設定の選択に利用できないことです。デフォルトのオプションは、ベイズ最適化を適用して有効なハイパーパラメータ値の領域を検索するもので、推奨されるアルゴリズムとなっています。このベイズ最適化の仕組みの詳細を知りたい方は、こちらのブログをご覧ください。
ジョブが開始されると、UI の [HYPERPARAMETER TUNING JOBS] タブでステータスを確認できるようになります。
終了すると、ジョブ名をクリックして調整のトライアル結果を見ることができます。
次のステップ
これで、Vertex Training でのハイパーパラメータ調整の使い方の基本がわかりました。最初から最後まで動作する例を試してみたい方は、こちらのチュートリアルをご覧ください。また、Vertex でのマルチワーカー トレーニングについて知りたい方は、こちらのチュートリアルをご覧ください。ぜひご自分でテストしてみましょう。
-デベロッパー アドボケイト Nikita Namjoshi
