~ベイズ最適化~ Optuna with SageMaker チューニングと評価を時短して効率化

26 min readJan 2, 2020

ここ数年で機械学習やディープラーニングのオープンソースライブラリは相当充実してきました（ChainerはPyTorchへ移行されるそうです）。あわせて機械学習やディープラーニングのアルゴリズムだけでなく、モデル選択やハイパーパラメータのチューニング自体もコモディティ化してきています。例えばオープンソースではauto-sklearnやTPOT、サービスとしてはDataRobotやGoogle Cloud AutoMLが挙げられ、これらのツール群はモデルの選択やパラメータの選択を自動化することが可能です。以下のグレー箇所は、TPOTで自動化される処理群です。[1]

今回試してみたことです。

OptunaとScikit-learnが使用できるSageMakerトレーニング用のコンテナ
このコンテナへパラメータとして最適化を試行する時間を指定実行する

https://www.semanticscholar.org/paper/Evaluation-of-a-Tree-based-Pipeline-Optimization-Olson-Bartley/3335c340c20609b4e6de481c9eaf67ecd6c960dc/figure/0

機械学習におけるこれらの処理の自動化はAutoML（Automated Machine Learning）として表現されます。AutoMLについてはこちらのスライドが非常に参考になりました。[2]

MLOpsにおいては機械学習やディープラーニングのデータクリーニングからモデル選択やハイパーパラメータのチューニングまで運用上に必要となる操作群をCI/CDパイプラインとして構築することが重要です。新しいデータに対するモデルの再構築やデプロイ、APIエンドポイントの更新などを手作業ではなく、自動化することでMLOpsとして円滑な開発・運用体制が構築できます。この機械学習のモデルの更新やAPIエンドポイントの作成は各種クラウドサービスを使用することで容易に実現できます。今回はAWSで提供されるAmazon SageMakerを使用してみました。

Amazon SageMakerは機械学習の開発・運用から推論サービスのデプロイまでを幅広くカバーするAWSのマネージドサービスで、SageMakerでは準備済みのコンテナで提供されているトレーニングアルゴリズムが使用できる他、Apache SparkやカスタムアルゴリズムをDockerイメージにまとめて、ECR（Elastic Conteiner Registry）から呼び出して使用することが可能となっています。基本的にDockerコンテナの知識が必要になります。[3]

https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/how-it-works-training.html

auto-sklearnやTPOTをSageMakerのトレーニングジョブを行うコンテナへインストールすることでモデル選択やハイパーパラメータのチューニングを自動化することができます。ただし一般的にモデルの選択やハイパーパラメーターのチューニングの探索は膨大な組み合わせの試行となるため、大きな計算資源と時間を要します。MLOpsにおいてモデルの再構築などを行うCI/CDパイプラインを組む場合に、計算資源および時間は限られてきます。つまり最適というより「ある程度最適」なモデルとハイパーパラメーターで妥協しても良いケースもあると思います。日次のデータから毎日モデル更新を行うことを考えた場合、この更新作業に膨大な時間はかけられません。例えば、朝の9時にモデル更新をかけて更新の完了が夕方では更新されたモデルを使用できるのは、夕方のモデルが更新された時点から翌日9時までとなってしまいます。AutoMLで自動化されたモデル選択やハイパーパラメーターのチューニングを行うauto-sklearnやTPOTでは一般的に探索の処理に非常に時間がかかります。TPOTのデフォルト設定（100 generations with 100 population size）は1万のパイプラインの試行ですが、これはグリッドサーチで1万個のモデル設定へ10分割の交差検証を適用し、10万個のモデルを評価することと同等だと考えられるようです。[4]

Using TPOT - TPOT

Automated machine learning (AutoML) takes a higher-level approach to machine learning than most practitioners are used…

epistasislab.github.io

仮想2コア、1GメモリのVPSでauto-sklearnやTPOTをirisデータセットで試したところ以下のような結果が得られました。auto-sklearnはメタ学習（Meta Learning）で効率的にMLフレームワークを絞り出し、ベイズ最適化によってモデルやパラメータの選択を行います。TPOTは遺伝的アルゴリズムを採用しています。最適化のアルゴリズムは確率論的な手法であり、複数のTPOTを走らせた場合に異なるパイプライン設定を推奨する場合があります。これはTPOTが時間不足のために収束しなかったか、または複数のパイプライン設定が同じデータセットで同等のパフォーマンスを実行することを意味するようです。irisデータセットのような軽量なものでも、それなりに時間がかかっていることが分かると思います。

まずauto-sklearnでは3600秒、つまり1時間ほどかかっています。使用したコードです。

Accuracy score 0.9736842105263158
Elapsed time: 3608.4015102386475[sec]

次はTPOTです。同じスコアが出力されているので、同じモデル・パラメーターが引き当てられているのかもしれません。1015秒、約17分かかっています。使用したコードです。

Accuracy score 0.9736842105263158
Elapsed time: 1015.0303757190704[sec]

※TPOTのオプティマイザーのパラメータは以下を使用、 generations をデフォルトの100から10へ引き下げています

pipeline_optimizer = TPOTClassifier(generations=10, population_size=100, cv=5,                                     random_state=42, verbosity=2)

本稿ではベイズ最適化を使用したハイパーパラメーターのチューニングのライブラリであるOptunaを時間指定することで実行し、ハイパーパラメーターチューニングされたモデルの構築をSageMakerで行います。Optunaでチューニングにかかる時間をコントロールすることでCI/CDのパイプラインを構築した際、ハイパーパラメーターのチューニングを実用的なレベルで行えるコンテナを作成します。 TPOTではTPOTのパラメータ設定はできるものの試行時間については制御ができません。与えられたデータセットやデータの累積によってモデル選定やパラメータチューニングにかかる時間も増加する可能性があります。例えば必ず5分間でトレーニングジョブやモデルの構築が完了するように時間内でハイパーパラメータのチューニングを終えたい動機が発生し得ます。

OptunaとScikit-learnが使用できるSageMakerトレーニング用のコンテナ
コンテナへパラメータとして最適化を試行する時間を指定実行する

Optunaとは

Optunaはベイズ最適化の一種であるTPE（Tree-structured Parzen Estimator）をデフォルトの最適化アルゴリズムとして持つハイパーパラメータのチューニングライブラリです。従来のグリッドサーチやランダムサーチとは異なり、過去の評価履歴から次の探索を決定するベイズ最適化を採用しています。PyData.Tokyo Meetupの講演資料がグッドリファレンスです。[5]

探索空間定義（Define-by-run）
枝刈り対応
分散最適化対応

PyData.Tokyo Meetup #21 講演資料「Optuna ハイパーパラメータ最適化フレームワーク」

Optunaはpipでインストール可能です。今回はSageMaker用のコンテナイメージに含めてしまいます。※ローカルに通常インストールする場合は以下

$ pip install optuna

Optuna付きコンテナの作成

ボストン住宅価格のデータセットのcsvファイルをデータとして取り、勾配ブースティングの回帰モデルを構築するSageMaker用のDockerコンテナを作成しました。このコンテナイメージは train による訓練およびモデルの構築、 serve による推論サービスのデプロイの両方をサポートします。パラメータとしてOptunaを実行する秒数を secondsとして与えられるようにします。デフォルトでは300秒として、jsonファイルで設定できるようにしておきます。公式サンプルのコードを改良して作成しました。[6]

awslabs/amazon-sagemaker-examples

This example shows how to package an algorithm for use with SageMaker. We have chosen a simple scikit-learn…

github.com

作成したDockerfileは以下です。optunaというディレクトリに train や serve など必要なアプリケーションコードを含めました。 train以外のコードはほぼサンプルのままです。

Necessary application codes inside the container

SageMaker Docker container with Optuna

requirements.txtにpip3でインストールする各種ライブラリを記載しました。推論サービスに必要なライブラリであるnginx/flask/gevent/gunicornや今回使用するOptunaなどがあります。 scikit-learnのバージョンをサンプルの 0.20.2 から 0.21.2 へ変更しています。

# requirements.txtnumpy==1.16.2
scipy==1.2.1
scikit-learn==0.21.2
pandas==0.25.3
optuna
flask
gunicorn
gevent

Optunaは評価のための目的関数を設定し、デフォルトで目的関数が最小化されるようにその最適解を求めます。公式サイトのチュートリアルでは試行回数100回で二次方程式の解を求める例が掲載されています。[7]

1.9487825780924659が求められ、実際の解である2に近い値が得られることが分かると思います。この例では目的関数の最小化、試行回数は100回という設定でOptunaによる最適化が行われています。

def objective(trial):
    x = trial.suggest_uniform('x', -10, 10)
    return (x - 2) ** 2study = optuna.create_study()
study.optimize(objective, n_trials=100)

本稿ではKFoldを使用した5分割交差検証によるR2スコアの平均を最大化するパラメータを求める目的関数としてコードを記載しました。デフォルトでは目的関数の最小化となりますが、 optuna.create_study へ direction='maximize' を渡すことで目的関数の返り値の最大化とすることができます。最小化の目的関数へ変更して回帰問題に対する評価値の1つであるRMSE（Root Mean Squared Error）などを使用しても良いです。RMSEの最小化は二乗誤差の最小化であるため、誤差に正規分布を仮定した場合の最尤推定と一致します。また study.optimize を呼び出す際に試行回数である n_trial= ではなく代わりに timeout=seconds を指定しています。 seconds は param_path で与えられるパスにあるjsonファイルに試行時間を秒数単位で設定します。/opt/ml/input/config/hyperparameters.jsonはハイパーパラメータをモデルに渡すためのパスですが、ここでOptunaの設定を渡しています。

# Use Optuna to find the hyperparameters with the highest score
print('Hyperparameter optimization starting.')
objective = Objective(X_train, y_train.values.ravel())
study = optuna.create_study(direction='maximize')
study.optimize(objective, timeout=seconds)
print("Optimized params: {}".format(study.best_params))

trainファイルは以下のようになりました。

SageMaker Docker train with Optuna

Dockerfileができたら任意の名前を付けてビルドしておきましょう。作成したコンテナ名を渡すことでシェルスクリプトによるモデル構築や推論サービスのテストがローカルで行えます。今回作ったコンテナイメージはDockerHubへアップロードしています。[8]

Docker Hub

Edit description

hub.docker.com

$ docker build -t optuna-sklearn-container .
$ docker images
REPOSITORY         TAG    IMAGE ID      CREATED             SIZE
optuna-sklearn-container latest b1cd2c6801af 24 minutes ago 1.29GB

ローカル環境でテスト

SageMaker用のコンテナイメージが作成できました。 local_test 配下にローカルテストを行うためのシェルスクリプトを配置してありますので試してみます（サンプルのスクリプトと同じです）。 ./train_local.sh optuna-sklearn-container でコンテナの train を実行することでモデルの訓練が進みます。シェルスクリプトの中で呼ばれているのは docker run -v $(pwd)/test_dir:/opt/ml --rm optuna-sklearn-container trainです。構築されたモデルがtest_dir/model配下に作成されていることが確認できます。ローカルのinput/config/hyperparameters.jsonには {"seconds": 300} として5分間をパラメータとして渡してあります。

$ ./train_local.sh optuna-sklearn-container
Starting the training.
X shape: (506,13)
y shape: (506,1)
Hyperparameter optimization starting.
[I 2020-01-01 06:47:35,320] Finished trial#0 resulted in value: 0.7267604802651177. Current best value is 0.7267604802651177 with parameters: {'n_estimators': 42, 'learning_rate': 0.9447497845654071}.
...[I 2020-01-01 06:52:33,147] Finished trial#667 resulted in value: 0.8742550361948627. Current best value is 0.8841786417221866 with parameters: {'n_estimators': 79, 'learning_rate': 0.13008465444573916}.
Optimized params: {'n_estimators': 79, 'learning_rate': 0.13008465444573916}
Evaluation for test data: 0.8157877127571357
Training completed.

300秒（5分間）で668回の試行が行われ、r2スコアの最大値は0.8841786417221866となりました。この最適化されたパラメータを使用してモデルを構築し、テストデータで評価したr2スコアが0.8157877127571357です。元のデータを訓練用とテスト用として8対2で分割しており、訓練用の8割のデータを使用して5分割の交差検証を用いたスコアの評価の相加平均値をOptunaでは採用しています。8割のデータ内の交差検証でスコアの評価によるハイパーパラメーターの最適化を行い、2割のデータに対して予測したものが Evaluation for test data として表示されているr2スコア値です。ボストン住宅価格のデータセットに対するr2スコアとしては比較的優秀なのではないでしょうか。

serve を実行する ./serve_local.sh optuna-sklearn-container をローカルで実行します。呼ばれているのはdocker run -v $(pwd)/test_dir:/opt/ml --rm optuna-sklearn-container serve です。

$ ./serve_local.sh optuna-sklearn-container
Starting the inference server with 2 workers.
[2020-01-01 06:55:20 +0000] [9] [INFO] Starting gunicorn 20.0.4
[2020-01-01 06:55:20 +0000] [9] [INFO] Listening at: unix:/tmp/gunicorn.sock (9)
[2020-01-01 06:55:20 +0000] [9] [INFO] Using worker: gevent
[2020-01-01 06:55:20 +0000] [13] [INFO] Booting worker with pid: 13
[2020-01-01 06:55:20 +0000] [14] [INFO] Booting worker with pid: 14
...# after sent the 10 records
Invoked with 10 records
172.17.0.1 - - [01/Jan/2020:07:14:38 +0000] "POST /invocations HTTP/1.1" 200 185 "-" "curl/7.47.0"

ローカルサーバに対してhttp:8080のポートで推論サービスをテストするには ./predict.sh ./payload.csvを実行します。curlでcsvデータを推論サービスへ渡して構築したモデルによる予測を出力できます。payload.csvはboston.csvの上から10行のデータをそのまま残したものです。元データのMEDVへかなり近い値が予測されていることが分かります。

$ ./predict.sh ./payload.csv
*   Trying 127.0.0.1...
* Connected to localhost (127.0.0.1) port 8080 (#0)
> POST /invocations HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/7.47.0
> Accept: */*
> Content-Type: text/csv
> Content-Length: 657
>
* upload completely sent off: 657 out of 657 bytes
< HTTP/1.1 200 OK
< Server: nginx/1.14.2
< Date: Wed, 01 Jan 2020 07:14:38 GMT
< Content-Type: text/csv; charset=utf-8
< Content-Length: 185
< Connection: keep-alive
<
25.17566973042558
21.631749140267193
34.32727090913451
34.17185574659824
35.031669916494195
24.77183077178688
20.962476186893326
22.250091728652645
16.441872773419494
18.74668242399064

Optunaの各trialの試行結果はDataFrame形式で取り出すことが可能です。csv形式へと変換しoutput_pathへファイル保存する処理を加えています。

hist_df = study.trials_dataframe()
hist_df.to_csv(os.path.join(output_path, 'optimization.csv'))

ローカルでSageMaker用のDockerコンテナが動作することを確認できました。このイメージでは python 3.7.5 で scikit-learn 0.21.2 を使用した勾配ブースティングによる訓練や推論が行えます。

ECR（Elastic Container Registry）へイメージをプッシュ

SageMakerのJupyter Notebookから作成したDockerコンテナを呼び出して使用するためには、コンテナイメージをECR（Elastic Container Registry）へアップロードする必要があります。サンプルのレポジトリでは build_and_push.sh として一連の流れをシェルスクリプトで行っています。

手動で実施する場合は以下の手順です。

AWS ECRにログインします。

$ aws ecr get-login --no-include-email
$ docker login -u AWS -p https://<account>.dkr.ecr.<region>.amazonaws.comLogin Succeeded

ECRにリポジトリを新規に作成します。

$ aws ecr create-repository --repository-name "optuna-sklearn-container"
{
    "repository": {
        "repositoryUri": "<account>.dkr.ecr.<region>.amazonaws.com/optuna-sklearn-container",
        "registryId": "<registry id>",
        "imageTagMutability": "MUTABLE",
        "repositoryArn": "arn:aws:ecr:<region>:<account>:repository/optuna-sklearn-container",
        "repositoryName": "optuna-sklearn-container",
        "createdAt": 1573965078.0
    }
}# Describe the existing repository
$ aws ecr describe-repositories

tagを付けて作成したECRのレポジトリへコンテナイメージをプッシュします。削除は aws ecr delete-repository で行えます。

$ docker tag optuna-sklearn-container <account>.dkr.ecr.<region>.amazonaws.com/optuna-sklearn-container
$ docker push <account>.dkr.ecr.<region>.amazonaws.com/sklearn-container# If deletion is needed, run this
$ aws ecr delete-repository --repository-name optuna-sklearn-container --force

ECRに使用したいコンテナイメージがプッシュできたので、次はSageMakerのノートブックインスタンスから実際にこのコンテナを使用して、モデルの訓練や、推論サービスのエンドポイントのデプロイを行ってみたいと思います。

SageMakerで試してみる

SageMaker Python SDKを使用してモデルをトレーニングします。Optunaでハイパーパラメータが最適化されたモデルが保存されるはずです。自前のコンテナイメージを使用する場合は sagemaker.estimator.Estimator のインスタンスからECRのリポジトリURI、トレーニングインスタンスなどを呼び出してください。[9]

/opt/ml/input/config/hyperparameters.jsonへ値を渡すには set_hyperparameters を使用します。

params = dict(seconds = 300)
clf.set_hyperparameters(**params)

ローカルで実行した場合と同じようにOptunaによるハイパーパラメータの最適化が走りました。トレーニングジョブ自体の実行時間は352秒となっています。Optunaの最適化は300秒と設定しましたが、MLインスタンスの準備などに若干時間がかかるためです。

2020-01-02 06:51:21 Completed - Training job completed
Training seconds: 352
Billable seconds: 352

構築されたモデルは指定したS3バケットの output というフォルダ配下に保存されます。モデルが保存されていることが確認できたので、SageMakerで推論サービスとしてデプロイしてみます。

deploy を呼び出してモデルの推論サービスをデプロイしました。

from sagemaker.predictor import csv_serializer
predictor = clf.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge", serializer=csv_serializer)

SageMakerのノートブックでローカルと同様にテストデータを渡して推論サービスを利用してみます。

SageMaker Container train with Optuna inference

以下のような結果となりました。それなりの精度で予測できているようです（元データを利用しているので当然の精度ですが、トレーニングジョブへ渡しているデータは全体の8割であり、これらのサンプルが含まれていたかは見ていません）。

Predicted values:
[25.1578253  22.42520892 33.77815577 34.55692806 35.69335112 25.32593544
 20.98115696 20.92998774 17.36359306 18.27929043]
test_y values:
[24.  21.6 34.7 33.4 36.2 28.7 22.9 27.1 16.5 18.9]

使用したコードとJupyter Notebookはgithub上に置いてあります。[10]

https://github.com/yuyasugano/sagemaker-optuna-container/blob/master/Scikit-learn_Estimator_Example_With_Optuna-Container.ipynb

https://github.com/yuyasugano/sagemaker-optuna-container

まとめ

Amazon SageMakerは機械学習モデルの構築、トレーニング、デプロイなどを提供するAWSのサービスである
AutoMLとしてモデルの選択からハイパーパラメータのチューニングまで自動化するソリューションが複数登場している、オープンソースではauto-sklearnやTPOT、商用ではDataRobotやGoogle Cloud AutoML等
Optunaによるハイパーパラメータのチューニングを行うコンテナを作成し、Optunaの実行時間をパラメータとして渡した
指定した実行時間の最適化でそれなりの精度を持つモデルの構築およびSageMakerによる推論サービスのデプロイができた