MLflow for Bayesian Experiment Tracking - The Databricks Blogの翻訳です。
この記事はベイジアン推論シリーズの3つ目の記事です([1]、[2])。ここでは、PythonパッケージのPyMC3を用いたベイジアンエクスペリメントを実行し、トラッキングするためにDatabricksのマネージドMLflowをどのように使うのかを説明します。バージョン管理機能、変数トラッキング機能によって、データサイエンスチーム間で共有可能かつ、システマティック、再現可能な実験MLパイプラインを構築することができます。MLflowによってトラッキングされるデータは、Databricks上のUI、APIを用いてアクセスすることができます。マネージドのMLflowサービスを使用していないデータサイエンティストは、エクスペリメントおよび関連づけられているデータに対してAPI経由でアクセスすることができます。Databricksにおいては、データ、モデルへのアクセスはMLflowが提供するアクセスコントロールリストによって管理されます。様々なフレームワークにモデルをデプロイすることで、モデルを容易に本格運用することができます。
ベイジアンエクスペリメントのトラッキング
MLflowは何をするのか？
MLflowはMLライフサイクルを管理するためのオープンソースのフレームワークです。MLflowはDatabricksのマネージドサービスあるいは、オープンソースのライブラリを用いてスタンドアローンデプロイメントとしてインストールすることができます。この記事では主にエクスペリメントのトラッキングを取り扱いますが、トレーニングしたモデルを集中管理リポジトリに格納する際、モデルデプロイメントの際にMLflowがどのように役立つのかに関しても説明します。トラッキングの文脈においては、MLflowは以下のものを格納することができます。
- メトリクス - 通常はデビアンスやRハットなどのモデルのパフォーマンスに関するメトリクスとなります。
- パラメーター - モデルやランを定義するのに使用した変数となります。ベイジアンの場合には、ハイパーパラメーターや事前分布パラメーター、超事前分布パラメーターとなります。これらの値は常に文字列として格納されることに注意してください。
- タグ - 特徴量を追加するコードのメジャーバージョンに関する情報のような、ランに関する情報を追跡するためのキーバリュー値となります。
- ノート - MLflowのUIで入力できるランに関する情報となります。ランの結果の定量的評価やシステマティックな実験で有用なツールとなります。
- アーティファクト - ファイル、画像など実験結果、副産物といったものを格納します。
オープンソースMLflowのストアのセットアップ
このセクションは、オープンソースのMLflowにのみ適用されます。DatabricksでホストされるMLflowではこれらは自動でケアされます。MLflowにはバックエンドストアとアーティファクトストアがあります。名前が示す通り、アーティファクトストアはモデルのランに関わる全ての(メタデータを含む)アーティファクトを保持し、それ以外はバックエンドストアに格納されます。ローカルでMLflowを実行している場合には、ファイルストレージあるいはデータベースのストアをバックエンドストアに設定することができます。以下に示すように、好きな場所でトラッキングサーバーを稼働させることができます。
mlflow server \
--backend-store-uri /mnt/persistent-disk \
--default-artifact-root s3://my-mlflow-bucket/ \
--host 0.0.0.0
その後で、上で設定したトラッキングサーバーを指定します。
mlflow.set_tracking_uri("http://YOUR-SERVER:4040")
ベイジアンエクスペリメントのトラッキングワークフロー
Databricksでは、あなたのために全てが管理されており、モデル開発ワークフローをスタートする際に必要となる設定の時間を最小化します。しかし、以下のステップは、マネージド、オープンソースのMLflowデプロイメントの両方で必要となります。MLflowは、エクスペリメントIDで識別されるエクスペリメントを作成します。それぞれのエクスペリメントは、ランIDで識別される一連のランから構成されます。それぞれのランは、ランごとに記録されたアーティファクトと関連づけられたパラメーターを保持します。ワークフローを構成するためのステップを以下に示します。
-
エクスペリメントのフォルダーへのパスを指定してエクスペリメントを作成すると、エクスペリメントIDを取得できます。ファイルや画像などのアーティファクトを格納するためのパスを指定することもできます。
-
上のステップで取得したエクスペリメントIDを指定してエクスペリメントをスタートします。PyMC3の推論コードは、コンテクストマネージャの配下になります。
-
コードとデータのバージョン管理を行うためにタグを使用します。
-
モデル、ランのパラメーターを記録します、特に事前、超事前分布パラメーター、サンプル数、チューニングのサンプル、尤度分布を記録します。
Python
mlflow.set_tags({"Version notes": "Full run", "Start date": covid_data.data_begin, "End date": covid_data.data_end})
mlflow.set_tags({"Version notes": "Full run",
"Start date": covid_data.data_begin,
"End date": covid_data.data_end})
nsamples = 2000
ntune = 2000
Hyperprior = {"Lambda mean": 0.75, "Lambda std": 2, "Mu mean": 0.75, "Mu std": 2}
Prior = {"Lambda std": 1.0, "Mu std": 1.0}
Likelihood = {"Name": "Normal", "Parameters": {"std": 0.01}}
prior_lam = pm.Lognormal('prior_lam', Hyperior['Lambda mean'], Hyperior['Lambda std'])
prior_mu = pm.Lognormal('prior_mu', Hyperprior['Mu mean'], Hyperprior['Mu std'])
prior_lam_std = pm.HalfNormal('prior_lam_std', Prior['Lambda std'])
prior_mu_std = pm.HalfNormal('prior_mu_std', Prior['Mu std'])
lam = pm.Lognormal('lambda', prior_lam , prior_lam_std, shape=2)
mu = pm.Lognormal('mu', prior_mu , prior_mu_std, shape=2)
mlflow.log_param("Hyperprior", Hyperprior)
mlflow.log_param("Prior", Prior)
mlflow.log_param("Samples", nsamples)
mlflow.log_param("Tuning samples", ntune)
mlflow.log_param("Likelihood", Likelihood)
1. モデルがサンプリングを完了すると、トレースに含まれる結果はlog_artifacts()メソッドを用いて保存されます。これは、それぞれのチェーンによって導かれたサンプルに関する全ての情報を含む`trace`というフォルダーとなります。PyMC3のsummary()メソッドをトレースオブジェクトに適用することで、トレース情報を要約することができます。トレースのサマリーは、MLflowのlog_text()メソッドを用いてJSON文字列オブジェクトとして保存できるデータフレームとなります。
```py:Python
trace_summary = az.summary(trace)
res = trace_summary.to_json(orient='index')
read_json = json.dumps(res)
mlflow.log_text(read_json, artifact_file = 'trace_summary')
pm.save_trace(trace, directory="trace", overwrite=True)
os.system('cp -R trace /dbfs/Users/USERNAME/mlflow')
mlflow.log_artifacts('/dbfs/Users/USERNAME/mlflow/trace', artifact_path = 'trace')
mlflow.end_run()
エクスペリメントの調査
エクスペリメントが完了すると、MLflowのUIで結果を調査するか、プログラムからランの情報を抽出することができます。例えば、現在のエクスペリメントIDが10618537である場合、エクスペリメントに関する情報を以下のように取得することができます。
client = mlflow.tracking.MlflowClient()
experiment = mlflow.get_experiment(10618537)
print(experiment)ed
<Experiment: artifact_location='dbfs:/databricks/mlflow-tracking/10618537', experiment_id='10618537', lifecycle_stage='active', name='/Users/USERNAME/mlflow_hier_bayesian', tags={'mlflow.experimentType': 'MLFLOW_EXPERIMENT',
'mlflow.ownerEmail': 'USEREMAIL',
'mlflow.ownerId': 'USERID'}>
エクスペリメントランの検索
エクスペリメントIDを知っているのであれば、以下のようにエクスペリメント内の全てのランを検索し、当該のランとして格納されているデータを取得することができます。
>>> mlflow.search_runs(experiment_ids = ["10618537"])
run = client.get_run('393287b22e59466dadbaea85052782a9')
print(run.data)
<RunData: metrics={}, params={'Hyperprior': "{'Lambda mean': 0.75, 'Lambda std': 2, 'Mu mean': 0.75, 'Mu "
"std': 2}",
'Prior': "{'Lambda std': 1.0, 'Mu std': 1.0}"}, tags={'End date': '4/1/20',
'Start date': '3/1/20',
'Version notes': 'Added artifacts',
'mlflow.databricks.cluster.id': '0512-212345-arbor51',
'mlflow.databricks.cluster.info': '{"cluster_name":"test","spark_version":"7.6.x-cpu-ml-scala2.12","node_type_id":"i3.xlarge","driver_node_type_id":"i3.xlarge","autotermination_minutes":120,"disk_spec":{},"num_workers":0}',
'mlflow.databricks.cluster.libraries': '{"installable":[{"pypi":{"package":"pymc3"}}],"redacted":[]}',
'mlflow.databricks.notebookID': '10564146',
'mlflow.databricks.notebookPath': '/Users/USERNAME/Hierarchical_downsampled_mlflow',
'mlflow.databricks.notebookRevisionID': '1626229409220',
'mlflow.databricks.webappURL': 'https://demo.cloud.databricks.com',
'mlflow.source.name': '/Users/USERNAME/Hierarchical_downsampled_mlflow',
'mlflow.source.type': 'NOTEBOOK',
'mlflow.user': 'USEREMAIL'}>
ランからアーティファクトへのアクセス
以下のように、このランに関連づけられているアーティファクトの一覧を表示することができます。ファイルごとのファイルサイズ、パスが表示されます。
f = client.list_artifacts(run_id)
print(f)
[<FileInfo: file_size=1808, is_dir=False, path='trace_summary'>,
<FileInfo: file_size=None, is_dir=True, path='trace'>]
MLflowは、それぞれのランごとにアーティファクトを管理しますが、UIやAPIを用いて参照、ダウンロードすることも可能です。以下の例では、事前のランからトレース情報とトレースサマリーをロードしています。
local_path_summary = client.download_artifacts(run_id, "trace_summary", './')
local_path_trace = client.download_artifacts(run_id, "trace", './')
print("Artifacts downloaded in: {}".format(local_path_summary))
print("Artifacts downloaded in: {}".format(local_path_trace))
with open(local_path_summary,'r') as f:
data = json.load(f)
trace_summary = pd.read_json(data)
with pm.Model() as model:
trace2 = pm.load_trace(local_path_trace)
data_load = az.from_pymc3(trace=trace2)
az.plot_posterior(trace2.get_values('R0')[:,0])
az.plot_posterior(trace2.get_values('R0')[:,1])
上を実行することで、トレースサマリーには以前と同じ情報が含まれていることに気づくことでしょう。アーティファクトファイルやトレースサマリーからロードしたパラメーターの推定値は、それらの分布が示すように、現在のモデルのパラメーターとなります。必要であれば、現在推定されている事後確率を将来のトレーニングサイクにおける事前確率とすることで、新規データに対するフィッティングを継続することができます。
まとめ
この記事では、PyMC3を用いたベイジアンエクスペリメントをシステマティックに実行するために、どのようにMLflowを使うのかを見てきました。MLflowによって提供されるロギング、トラッキング機能は、DatabricksのマネージドMLflow、あるいはオープンソースMLflowのユーザーであればAPIを通じて利用することができます。モデル、モデルのサマリーはアーティファクトとして保存され、共有したり、後でPyMC3に再読み込みすることができます。
詳細に関しては、添付のノートブックをチェックしてください。
マネージドMLflowやベイジアンエクスペリメントに関しては、ノートブックをチェックし得ください。私のCouseraのコースでもベイジアン推論を学ぶことができます。