【MLflow】Databrciks 環境で実験トラッキング – Python-

はじめに

Databricks環境下でオープンソースプラットフォームである MLflow を利用します。
実験データの統合・トラッキング方法を解説します。
また、本記事は Python3 を想定しています。

この記事の概要

  • Databricks 環境で MLfow を使うことできること
  • MLflow の使い方
  • サンプルデータを用いた実験の情報をノートブックに統合・トラッキング方法

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線形回帰モデルを使います。
調整パラメータとしてalpha、l1_ratioがあります。

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

セットアップ

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

結果処理部分の実装

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

実験処理部分の実装

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

実験

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

出力結果は次のようになりました。

https___qiita-image-store.s3.ap-northeast-1.amazonaws.com_0_618791_83f4d625-03b3-a99b-2d41-582a98b8ba97.png

画像を出力してみます。
https___qiita-image-store.s3.ap-northeast-1.amazonaws.com_0_618791_101943a6-9207-af62-380a-ca11bd234cc1.png

いくつかパラメータを変えて実験してみます。(合計4パターンほどやってみました)
実験が終わったら右上のあたりに[Runs]と書かれた部分を押します。

https___qiita-image-store.s3.ap-northeast-1.amazonaws.com_0_618791_6dd4d923-3f5c-5847-8325-037740512075.png

設定したパラメータの値と出力されたメトリクスが実験ごとに記録されているのが確認できます。

https___qiita-image-store.s3.ap-northeast-1.amazonaws.com_0_618791_2ab43d48-afff-2de2-48ff-429a7812cbee.png

おわりに

今回はノートブック上でモデルのトレーニング結果をノートブックに統合・トラッキングする方法について紹介しました。

Databricks では UI 上でこちらの実験データを比較ができるようになっています。
続編として、UI 上のモデル管理について書きたいと思います。

次回もお楽しみに!