Azure Machine Learning と MLflow と使って最速で実験管理を行う方法


はじめに

こんにちは、修士時代に機械学習をしていた折、「model_adam_lr_0001_lstm_3_layer_e_512_h_1024.model」みたいな名前を付けてモデル管理をしていたイトーと言います。

今回は既にある Jupyter ノートブックに最小限の手間で実験管理を組み込んで上記のようなヤバい名前のファイルを手動でつける事態を回避し、さらには論文やらの締め切り間際に焦りすぎて「model_adam_lr_0001_lstm_3_layer_e_512_h_1024 - コピー (最新版) (こっちが最新) version2 コピー (最終版) .model」みたいな事態になって完全にお終いにならないようにするための方法を提案します。

私はマイクロソフトで Azure を取り扱うアーキテクトですので、今回は実験管理ツールとして Azure Machine Learning を中核に据えます。が、ほとんど GUI が目当てみたいなもので大したことはしません。実質 MLflow の説明です。

手順

Azure Machine Learning のデプロイ

Azure で「機械学習」と検索欄に打ち込んで Azure Machine Learning を探します。

リソース作成ページに進んだらリソースグループやらストレージアカウントやら新規作成とあるものは全部新規作成を押して、適当な名前を付けてください。

全部名前だけ打ち込んで、あとはデフォルトのままで良いです。

一通り名前を付け終わったら最後にワークスペース名を付けて、左下の「確認と作成」→「作成」とボタンを押してリソースを作ります。

一刻も早く実験管理をしたいので細かい話は割愛します。

実験管理準備

pip install azureml-sdk
pip install mlflow
pip install azureml-mlflow

普段機械学習実験に使ってるカーネルに上記3つのライブラリを入れます。

入れ終わる頃には Azure Machine Learning のデプロイが終わってると思いますので、連携に必要な情報を手元に落としておきます。

config.json をダウンロードしてください。後に必要となる情報が全部記載されています。

実験管理

まず以下のコードを実験用のノートブックの一番最初のセルに挿入します。(コードの中身が分かる人は適宜好きな位置に挿入してください)

from azureml.core import Workspace
import mlflow

subscription_id = "<id>" # 環境ごとに変更
resource_group = '<rg>' # 環境ごとに変更
workspace_name = '<ws>' # 環境ごとに変更

ws = Workspace(
        workspace_name=workspace_name,
        subscription_id=subscription_id,
        resource_group=resource_group,
    )

mlflow_uri = ws.get_mlflow_tracking_uri()
mlflow.set_tracking_uri(mlflow_uri)

experiment_name = "<experiment_name>" # 好きな名前を指定
mlflow.set_experiment(experiment_name)

run_name = f'run_{int(time.time())}' # 実行ごとにかぶらないようにしつつ好きな名前を指定
run = mlflow.start_run(run_name=run_name)

「環境ごとに変更」のコメントを入れている部分は config.json の中身を見ながら値を書き入れておいてください。

セルを実行すると Output としてログインを要求するメッセージが出てくるので、 Azure で使用している Microsoft アカウントでログインしましょう。ログインが完了するとセルの実行が完了します。

続いてノートブックの最後のセルに以下のコードを挿入します。

mlflow.end_run()

続いて以下記載の関数を使ってメトリック等の記録を取ります。

# パラメーターを記録
mlflow.log_param("param_name", param)

# 辞書でパラメーターを定義している場合、丸ごと辞書を記録する
mlflow.log_params(dictionary)

# メトリックを記録
mlflow.log_metric("metric_name", metric_value)

# 辞書でメトリックをまとめている場合、丸ごと辞書を記録する
mlflow.log_metrics(dictionary)

# モデルなど、出力したファイルを記録する
mlflow.log_artifact(local_path, remote_path)

以上です。これでとりあえず最低限度実験の実行ごとに記録が取られる状態となりました。

「model_adam_lr_0001_lstm_3_layer_e_512_h_1024.model」とは金輪際縁を切れますね。今すぐ切りましょう。

確認

実験記録は Azure Machine Learning Studio から確認できます。

Studio の URL は Azure の作成したリソースの詳細ページに記載がありますのでそちらをお使いください。

image.png

大変よろしいですね。

モデル登録

mlflow.log_artifact()で記録したファイルをモデルとして別途登録・管理することが可能です。

model_name = "<model_name>"
mlflow.register_model(run.info.artifact_uri+"/"+model_path, model_name)

mlflow.end_run() の後に適宜実行しておきましょう。

良いモデルができたらそのときだけ登録するといった運用もありです。これでモデルを無くすことは2度となくなりますね。

そもそもモデルを毎回保存せずグラフだけ出力して満足している民は今すぐただちに可及的速やかにASAPでモデルを保存する癖をつけましょう。やらないといつかめちゃくちゃ後悔します。経験者は語るのです。

より簡単に

もし以下に記載されいてるフレームワークを利用している場合、mlflow.log_metricしなくてもよしなにログを取得してくれる Automatic Logging 機能が使用できるかもしれません。

サンプルコード

Github にサンプルコードを公開しています。MIT ライセンスのもとご自由にお使いください。

ちょっとだけ解説

クラウドサービス

実験管理ツールとして有名な OSS である MLflow というソフトウェアがあります。上記コードではほとんどこの MLflow を使用しています。

Azure Machine Learning は MLflow の互換エンドポイントを提供しており、今回はこの互換エンドポイント目当てに Azure Machine Learning をデプロイしています。要するに MLflow-as-a-Service が欲しかったということです。

あえてローカルに MLflow サーバーを立てずにクラウドサービスを利用している理由は以下3点です。

  • ローカルでサーバーを作るとローカルマシンが壊れた時に実験記録が失われて辛い
  • MLflow サーバーのセットアップが少々手間がかかる (DB のセット等)
  • MLflow では認証を提供できず、セキュリティに不安が残る

Azure Machine Learning を使うことでこの3点を解決できます。

なおこの使い方の場合、かかるコストは Storage にデータを保存する少々のコストのみとなります。Azure Machine Learning はマネージドで内部に構築できる計算リソースや依存リソース(ストレージ等)の課金のみです。

学生の場合は以下無料アカウントクレジットを使えたりしますが、このクレジットの金額を超えることはまずないかと思います。(GPT-2 の実験を繰り返しモデルを保存しまくるとか実験を5億回ぐらいやるといった限界事例の場合話は別です)

MLflow の実験管理

Experiment に複数の Run が所属するという形になっています。 (これは Azure Machine Learning でも同様で、それゆえに互換エンドポイントが提供できています)

experiment_name = "<experiment_name>" # 好きな名前を指定
mlflow.set_experiment(experiment_name)

この部分で実験を指定しています。同じ Experiment として管理したい場合はここを共通に、別の Experiment として管理したい場合はかぶらないように別の名前をつけましょう。

続いて Run ですが、以下のコードで Run の生成を行っています。

run_name = f'run_{int(time.time())}' # 実行ごとにかぶらないようにしつつ好きな名前を指定
run = mlflow.start_run(run_name=run_name)

mlflow.start_run()を実行したときから最後にmlflow.end_run()するまでの間、mlflows.log_metric()などを実行可能です。さらに start から end までの間の時間が記録されます。できるだけ実験本体に近い部分で start_run した方が正確な実験時間を計測可能ですが、あんまり気にしなくてもいいと思います。

その他にも便利な関数がありますが、詳細はこちらドキュメントで確認してください。

おわりに

Azure Machine Learning を MLflow-as-a-Service として取り扱い、最速で実験記録を取る方法を提案しました。

  1. Azure ML をデプロイする
  2. セルを2つばかり挿入して MLflow の記録準備を整える
  3. 実験記録用のコードを適宜挿入する

以上となります。

ちゃんと実験管理して「model_adam_lr_0001_lstm_3_layer_e_512_h_1024 - コピー (最新版) (こっちが最新) version2 コピー (最終版) .model」から脱却し、効率的に機械学習を進めていきましょう。