Amazon SageMaker で Optuna を用いたハイパーパラメータ最適化を実装する

Amazon SageMaker はお客様の機械学習のワークロードにおいて様々な選択肢を提供します。深層学習フレームワークの選択肢として2018年の AWS Summit Tokyo で発表された Chainer 対応はその一つです。Chainer は株式会社Preferred Networks により開発された深層学習フレームワークで、計算時に動的にグラフを生成する define-by-run の考え方 (imperative な実行とも呼ばれます) を世界に先駆けて取り入れました。株式会社Preferred Networks はこの Chainer とは独立に、同じく define-by-run の思想に基づいたハイパーパラメータの最適化 (HPO) のための Optuna を2018年12月に発表しました。本稿では、AWS が提供する SageMaker 上で Optuna を用いた HPO を行う方法とアーキテクチャについてご紹介します。

導入

SageMaker が提供する HPO の選択肢

Amazon SageMaker は、TensorFlow, Apache MXNet, PyTorch, Chainer, scikit-learn, Horovod, Keras, Gluon などのフレームワーク・インターフェースに対応し、すべての開発者とデータサイエンティストに機械学習モデルの構築・学習・デプロイ手段を提供する AWS のサービスです。SageMaker はマネージド型の Jupyter Notebook/Lab 環境、学習・デプロイ用のコンテナ環境などと合わせて、標準でベイズ最適化を用いた HPO の機能を提供しています。

SageMaker HPO について書かれた AWS ブログで説明されているように、SageMaker の HPO では Estimator を初期化した上で、以下のようにパラメータ空間を定義します。

from sagemaker.tuner import HyperparameterTuner, IntegerParameter, CategoricalParameter, ContinuousParameter
hyperparameter_ranges = {'optimizer': CategoricalParameter(['sgd', 'Adam']),
                         'learning_rate': ContinuousParameter(0.01, 0.2),
                         'num_epoch': IntegerParameter(10, 50)}

ログに出力されたメトリクスを正規表現によりパースし、HyperparameterTuner クラスを用いて HPO を実行します。

objective_metric_name = 'Validation-accuracy'
metric_definitions = [{'Name': 'Validation-accuracy',
                       'Regex': 'Validation-accuracy=([0-9\\.]+)'}]

tuner = HyperparameterTuner(estimator,
                            objective_metric_name,
                            hyperparameter_ranges,
                            metric_definitions,
                            max_jobs=9,
                            max_parallel_jobs=3)
tuner.fit({'train': train_data_location, 'test': test_data_location})

このように HPO を実行する際の最大ジョブ数と並列ジョブ数を指定し、非常に簡単にベイズ最適化を用いた HPO を行うことができます。

Optuna の概要・使い方

Optuna では define-by-run による HPO の記述ができます。例えばニューラルネットワークの層の数を最適化の対象とする場合、選択されるレイヤー数に応じたネットワーク構造の分岐的な探索など、複雑なパラメータ空間を予め定義するのは少し手間がかかります。このような場合、define-by-run のインターフェースをもつ Optuna では以下のようにモデル定義時に HPO のパラメータ空間を同時に定義することでより直感的な記述を可能にします。なお、この関数に渡される trial は Optuna により定義される引数で、これを用いてパラメータを取得します。

def create_model(trial):
    # We optimize the numbers of layers and their units.
    n_layers = trial.suggest_int('n_layers', 1, 3)

    layers = []
    for i in range(n_layers):
        n_units = int(trial.suggest_loguniform('n_units_l{}'.format(i), 4, 128))
        layers.append(L.Linear(None, n_units))
        layers.append(F.relu)
    layers.append(L.Linear(None, 10))

    return chainer.Sequential(*layers)

Optuna を用いるには学習スクリプトの中で、最小 (あるいは最大) 化したい値を返すよう、目的関数 objective() を下記のように定義します。

def objective(trial):
    # Model and optimizer
    model = L.Classifier(create_model(trial))
    optimizer = create_optimizer(trial, model)

    train_iter = chainer.iterators.SerialIterator(train, BATCHSIZE)
    test_iter = chainer.iterators.SerialIterator(test, BATCHSIZE, repeat=False, shuffle=False)

    # Trainer
    updater = chainer.training.StandardUpdater(train_iter, optimizer)
    trainer = chainer.training.Trainer(updater, (EPOCH, 'epoch'))
    trainer.extend(chainer.training.extensions.Evaluator(test_iter, model))
    log_report_extension = chainer.training.extensions.LogReport(log_name=None)
    trainer.extend(chainer.training.extensions.PrintReport(
        ['epoch', 'main/loss', 'validation/main/loss',
         'main/accuracy', 'validation/main/accuracy']))
    trainer.extend(log_report_extension)

    # Run!
    trainer.run()

    # Set the user attributes such as loss and accuracy for train and validation sets with SageMaker training job name.
    log_last = log_report_extension.log[-1]
    for key, value in log_last.items():
        trial.set_user_attr(key, value)

    trial.set_user_attr('job_name', args.training_env['job_name'])
    serializers.save_npz(os.path.join('/tmp', 'model_{}.npz'.format(trial.number)), model)

    # Return the validation accuracy
    return log_report_extension.log[-1]['validation/main/accuracy']

ここで、例えば create_model(trial) の中では n_layers = trial.suggest_int('n_layers', 1, 3) のような記述により、試行ごとにパラメータが選択されます。同様に、create_optimizer(trial, model) の中でも最適化手法や学習率などを HPO の対象として定義しています。

こうして定義された目的関数を以下のように呼び出すことで HPO が実行されます。

study = optuna.study.create_study(storage=db, study_name=study_name, direction='maximize')
study.optimize(objective, n_trials=100)

ここで study は HPO ジョブの単位で、指定したデータベース内に保存されます。今回は以下で説明するように Amazon Aurora 内のデータベースを指定します。

SageMaker で Optuna を利用するための構成

Optuna は pip install optuna により簡単にインストールすることが可能です。Optuna はデフォルトではインメモリに履歴を持ちますが、履歴を永続化させたい場合はリレーショナルデータベースをバックエンドとして利用します。本ブログでは複数ノードにまたがる並列最適化を行う場合を想定して、AWS 上のマネージドデータベースサービスである Amazon Aurora MySQL を用いた構成を紹介します。Aurora は、AWS 上の他のネットワークと論理的に分けられた、お使いの AWS アカウント専用の仮想ネットワーク環境である Amazon Virtual Private Cloud (VPC) 内に立ち上がります。なお、データベースは外部から直接のインターネット接続を必要としないため、VPC 内のプライベートサブネットに配置します。

モデル開発のための Jupyter 環境である SageMaker ノートブックインスタンスも、これと同一の VPC 内に起動します。開発時にはインターネット経由で Jupyter 環境に接続するためこちらは Public Subnet に配置し、インスタンスの仮想ファイアウォールであるセキュリティグループとネットワークトラフィック経路を定義するルートテーブルを適切に設定し Aurora への接続を許可します。また、学習時のコンテナも同じく VPC 内に起動させ、ライブラリのインストールのためにインターネットへ接続が必要なので NAT Gateway も作成しています。

環境を AWS 上で簡単に立ち上げて試せるよう、GitHub レポジトリに AWS CloudFormation テンプレートとサンプルコードを置きました。「Stack の作成」を押し、AWS マネージメントコンソールから CloudFormation スタックを作成すると、数分で以下の構成が立ち上がります (CREATE_COMPLETE と表示されていることを確認して下さい)。なお、この CloudFormation スタックを立ち上げると課金が発生するため、不要になった CloudFormation スタックは削除するなどご注意下さい。

今回、データベース認証情報を安全に SageMaker に渡すため、AWS Secrets Manager にデータベースユーザー名とパスワードをシークレットとして保存しています。

SageMaker 上での実行

それでは実際に SageMaker ノートブックインスタンスに接続して試してみましょう。AWS マネージメントコンソールからリダイレクトして Jupyter 環境を開きます。GitHub からクローンした chainer_simple.ipynb サンプルコードを開いて下さい。カーネルは conda_chainer_p36 を選択します。なお、Optuna は Chainer とは独立したライブラリなので、他のフレームワーク (TensorFlow, MXNet, PyTorch など) やライブラリ (XGBoost, scikit-learn など) の HPO にも使えますが、ここでは Chainer を例として説明します。Python 側には Optuna と合わせて予め MySQL のコネクタをインストールしておきます (pip install mysql-connector-python)。

学習の準備

SageMaker では学習と推論エンドポイントのホスティングに Docker コンテナを使うよう設計されています。今回は追加で Optuna を入れるだけなので、既存の SageMaker Chainer コンテナに requirements.txt を使用して環境をセットアップします。以下のようなバージョンを指定し、コンテナ起動時にインストールすることが可能です。本稿のスコープからは外れるため詳細の説明は割愛しますが、予めこれら追加ライブラリを含む学習・推論用の SageMaker コンテナを作成することでコンテナ起動時のオーバーヘッドを削減することができます。

mysql-connector-python >= 8.0.15
optuna >= 0.9.0

SageMaker では、エントリーポイントとして指定した Python スクリプト (ここでは src/chainer_simple.py) が学習時に実行されます。通常の学習用コードから SageMaker 上で実行できる Python スクリプトへの書き換えの変更点はあまり多くありません。main guard の中で、学習ジョブ実行時に渡されたハイパーパラメータをパースし、S3 から取得したデータをロードし、学習後はモデルを保存します。具体的な書き方は Keras を例にしたブログを参考にしてください。

学習ジョブの実行

CloudFormation の Outputs から必要なパラメータ (学習コンテナを起動するプライベートサブネットの ID とセキュリティグループ、Aurora のエンドポイントと SecretsManager のシークレット名) を取得します。データは予め SageMaker Python SDK を用いて S3 にアップロードします。それでは、先ほど用意した学習スクリプトを SageMaker ノートブックインスタンスから呼び出します。その際に、Chainer のバージョン、実行のための IAM ロール、サブネット ID、セキュリティグループ、インスタンスの種類と数などを指定して Chainer Estimator クラスを初期化し、fit() メソッドでコンテナを起動します。サンプルノートブックではここで以下のように複数の学習ジョブを並列実行しています。

from sagemaker.chainer.estimator import Chainer

chainer_estimator = Chainer(entry_point='chainer_simple.py',
                            source_dir="src",
                            framework_version='5.0.0',
                            role=role,
                            sagemaker_session=sagemaker_session,
                            subnets=subnets,
                            security_group_ids=security_group_ids,
                            train_instance_count=1,
                            train_instance_type='ml.c5.xlarge',
                            hyperparameters={
                                'host': host,
                                'db-name': db_name,
                                'db-secret': secret_name,
                                'study-name': study_name,
                                'n-trials': 25,
                                'region-name': region_name
                            })

# HPO in parallel
max_parallel_jobs = 2
for j in range(max_parallel_jobs-1):
    chainer_estimator.fit({'train': train_input, 'test': test_input}, wait=False)
chainer_estimator.fit({'train': train_input, 'test': test_input})

学習中のコンソール出力は Jupyter Notebook 上に表示されます。モデルの保存は objective() 関数の中で毎回行われますが、各学習ジョブの中で最適だった trial のみが model_dir に移動され S3 に保存されるようになっています。なお、実行後の結果の取り出しを容易にするため、 trial のユーザー属性に SageMaker 学習ジョブ ID を記録しています。これにより、推論エンドポイントのデプロイ時に最良のモデルが保存されている S3 のオブジェクトキーを簡単に得ることができます。

学習ジョブが終了したら結果を可視化してみましょう。以下のように pandas.DataFrame として結果を取り出すことが可能です。

study = optuna.study.load_study(study_name=study_name, storage=db)
df = study.trials_dataframe()

ax = df['user_attrs']['validation/main/accuracy'].plot()
ax.set_xlabel('Number of trials')
ax.set_ylabel('Validation accuracy')

試行を重ねることでより良いハイパーパラメータを探索し、精度が向上していることを確認できます。

推論エンドポイントの作成

それでは最後にモデルをデプロイして推論結果を確認してみましょう。これまで実行した study の中で最良のモデルを取得し、その結果が保存されている S3 バケット/オブジェクトを先ほど記録した SageMaker ジョブ ID を用いて指定します。モデルの読み込みは chainer_simple.py の中で記述されています。インスタンスタイプと数を指定し deploy() メソッドを呼び出すと、数分でエンドポイントが立ち上がります。

from sagemaker.chainer import ChainerModel

study = optuna.study.Study(study_name=study_name, storage=db)

best_model_data = os.path.join(chainer_estimator.output_path, study.best_trial.user_attrs['job_name'], 'output/model.tar.gz')
best_model = ChainerModel(model_data=best_model_data,
                          role=role,
                          entry_point='chainer_simple.py',
                          source_dir="src")

predictor = best_model.deploy(instance_type="ml.m4.xlarge", initial_instance_count=1)

試しに5枚の画像をエンドポイントに送ると推論結果が表示されます。

import random
import matplotlib.pyplot as plt

num_samples = 5
indices = random.sample(range(test_data.shape[0] - 1), num_samples)
images, labels = test_data[indices], test_labels[indices]

prediction = predictor.predict(images)
predicted_label = prediction.argmax(axis=1)
print('The predicted labels are: {}'.format(predicted_label))

以下のような結果が出力されれば成功です。

The predicted labels are: [9 7 8 6 1]

最後に、使わなくなったエンドポイントは意図しない課金を防ぐために削除しておきましょう。

predictor.delete_endpoint()

また、不要であれば CloudFormation Stack ごと削除しましょう。AWS マネージメントコンソールから実行することができます。

まとめ

本稿では Amazon SageMaker 上で Optuna を使ったハイパーパラメータの最適化 (HPO) を行うため、現時点で提供される AWS サービスと Optuna のバージョンで必要な構成と実行方法について説明しました。なお今回は Amazon Aurora MySQL を使った方法を紹介しましたが、Study 用のデータベースに RDS を使って実装することもできます。このように、SageMaker と AWS のサービスを組み合わせることで、様々なワークロードを実行することが可能です。他にも便利なユースケースや新しいアイディアをお持ちであれば、ぜひ我々ソリューションアーキテクトに教えて下さい。

著者について

hariby

針原佳貴 (Yoshitaka Haribara, Ph.D.) は AWS のソリューションアーキテクトです。趣味はドラムで休日は友達とバンドをして過ごしています。好きなサービスは Amazon SageMaker です。

宇都宮聖子 (Shoko Utsunomiya, Ph.D.) は AWS の機械学習ソリューションアーキテクトです。自動車・ヘルスケア・ゲームなど様々な領域で、お客様の機械学習プロジェクト支援やソリューション開拓に従事しています。好きなサービスは Amazon SageMaker です。

Amazon Web Services ブログ