DatabricksでMLflowを使う① – ノートブック上での実験トラッキング –

はじめに

本記事では、機械学習モデルのライフサイクル管理を行うオープンソースプラットフォームであるMLflowをDatabricksの環境下でトラッキングする方法について書きます。(Python3を想定しています)

Databricksのマネージド型MLflow

オープンソースであるMLflowではあらゆるMLライブラリや言語、デプロイメントツールで動作するように設計されていますが、実験のトラッキング用のサーバーを自前で用意する必要があります。

Databricksの環境下では、MLflowはマネージド型のサービスとして使うことができるのでトラッキング用のサーバーを別途用意する必要がありません。また、実験のトラッキング情報をノートブックに統合して管理することもできます。

今回は実験の情報をノートブックに統合してトラッキングする方法について書きます。

MLflowの使い方

Databricks上で使っているクラスターがRuntime MLの場合は最初から入っていますが、それ以外の場合はMLflowをインストールする必要があります。

上記コマンドでインストールできます。そしてインポートします。

MLflowではトラッキング開始のモジュールを呼び出してトラッキングを開始し、実験のパラメータやログを記録するモジュールで記録し、トラッキング終了のモジュールで一つの実験トラッキングが終了するという形になります。

終了し忘れ防止のためにもwithを使うのがいいと思います。
実装部分のコードのイメージとしては以下のようになります。

パラメータ以外にもモデルの記録や実験中に出力した画像などもトラッキング先に保存することができます。
モデルの記録には、該当するライブラリのインストールが別途必要になります。
例)scikit-learnのモデル → mlflow.sklearn

実装

実際に実装してノートブック上でトラッキングしてみます。

サンプルデータセットと使用モデル

scikit-learnの糖尿病データセットを使います。
カラムの説明などはこちらにあります。
https://scikit-learn.org/stable/datasets/index.html#diabetes-dataset

今回のモデル作成にはElasticNet線形回帰モデルを使います。
調整パラメータとしてalphal1_ratioがあります。

ElasticNetについてはこちらの説明がわかりやすかったです。
https://aizine.ai/ridge-lasso-elasticnet/

セットアップ

各種ライブラリのインポートとサンプルデータセットの読み込み、データフレームの作成を行います。

結果処理部分の実装

ElasticNetで回帰モデルを作ったときの説明変数の各係数をプロットして、それを画像としてドライバノードに保存する処理を定義しています。

実験処理部分の実装

alphal1_ratioを指定してモデルのトレーニングを行います。上で定義したplot_enet_descent_pathを呼び出してトラッキング先にログや画像を保存します。

実験

調整パラメータを与えて実験を行います。

出力結果は次のようになりました。
スクリーンショット 2020-08-25 17.00.47.png

画像を出力してみます。
スクリーンショット 2020-08-25 17.00.56.png

いくつかパラメータを変えて実験してみます。(合計4パターンほどやってみました)

実験が終わったら右上のあたりに[Runs]と書かれた部分を押します。
スクリーンショット 2020-08-25 17.08.27.png

設定したパラメータの値と出力されたメトリクスが実験ごとに記録されているのが確認できます。
スクリーンショット 2020-08-25 17.10.19.png

おわりに

今回はノートブック上でのモデルのトレーニングの結果をノートブックに統合するしてトラッキングする方法について書きました。
DatabricksではUI上でこちらの実験データを比較することができるようになっています。
続編として、UI上のモデル管理について書きたいと思います。