MLFlowでLightGBMの学習結果をtrackingしてみる

結構前にMLFlowをいろいろ触ってみていたんですが、最近全然触っていなかったので色々見てみました。

www.nogawanogawa.com

前に自分が触っていたときよりだいぶ使いやすくなってたので、今回は最近の自分の用途に合わせて改めてMLFlowを使ってみます。

例題：LightGBMで回帰問題を解く
- Automatic Logging
- LGBMの木の構造を保存する
コードの残骸
参考文献
感想

例題：LightGBMで回帰問題を解く

今回はLightGBMを使用して回帰問題を解くコードがすでにある状況について考えることにするので、特にLGBMのタスク自体については考えないことにします。

なんの問題を解くのかはおいておいて、多分きっとこんな感じでコードが書かれているとします。

import lightgbm as lgb
import sklearn.preprocessing as sp
from sklearn.model_selection import train_test_split
import dataloader, preprpcessor


class LGBMRegressor:
    def __init__(self):
        pass

    def train(self, X_train, X_valid, y_train, y_valid):

        train_data = lgb.Dataset(X_train, label=y_train)
        validation_data = lgb.Dataset(X_valid, label=y_valid, reference=train_data)

        params = {
            'objective': 'regression',      
            'metric': "rmse",               
            'learning_rate': 0.05,
            'num_leaves': 21,               
            'min_data_in_leaf': 3,    
            'early_stopping':100,    
            'num_iteration': 1000           
            }

        self.model = lgb.train(params,
               train_set=train_data,
               valid_sets=validation_data,
               verbose_eval=50)
        
        viz = dtreeviz(self.model,
                    x_data = X_valid,
                    y_data = y_valid,
                    target_name = 'rating',
                    feature_names = X_train.columns.tolist(),
                    tree_index = 0)

        filename = 'lgb_tree.svg'
        viz.save(filename)
        
    def predict(self):
        pass

if __name__ == '__main__':
    df = dataloader.data()
    user = dataloader.user()
    item = dataloader.item()

    df = preprpcessor.preprocess(df, user, item)

    y = df['rating']
    X = df.drop(['rating'], axis=1)

    X_train, X_valid, y_train, y_valid = train_test_split(X, y,
                                                    test_size=0.2,
                                                    random_state=1,
                                                    stratify=y)

    lgb_regressor = LGBMRegressor()
    lgb_regressor.train(X_train, X_valid, y_train, y_valid )

この状態からMLFlowで実験結果をトラッキングすることを考えたいと思います。

Automatic Logging

MLflowには、メジャーなライブラリにAutomatic Loggingという”いい感じに実験ログを記録してくれる君”が用意されています。

www.mlflow.org

これを使うと、実はLightGBMなどのよく使われるライブラリなんかは1行追加するだけでトラッキングできちゃったりします。便利ですね。

実際に追加するのは下記。

import lightgbm as lgb
import sklearn.preprocessing as sp
from dtreeviz.trees import *
import mlflow # ←ここ
from sklearn.model_selection import train_test_split
import dataloader, preprpcessor


class LGBMRegressor:
    def __init__(self):
        pass

    def train(self, X_train, X_valid, y_train, y_valid):
        mlflow.lightgbm.autolog() # ←ここ

        train_data = lgb.Dataset(X_train, label=y_train)
        validation_data = lgb.Dataset(X_valid, label=y_valid, reference=train_data)

        params = {
            'objective': 'regression',      
            'metric': "rmse",               
            'learning_rate': 0.05,
            'num_leaves': 21,               
            'min_data_in_leaf': 3,    
            'early_stopping':100,    
            'num_iteration': 1000           
            }

        self.model = lgb.train(params,
               train_set=train_data,
               valid_sets=validation_data,
               verbose_eval=50)
        
        viz = dtreeviz(self.model,
                    x_data = X_valid,
                    y_data = y_valid,
                    target_name = 'rating',
                    feature_names = X_train.columns.tolist(),
                    tree_index = 0)

        filename = 'lgb_tree.svg'
        viz.save(filename)
        
    def predict(self):
        pass

if __name__ == '__main__':
    df = dataloader.data()
    user = dataloader.user()
    item = dataloader.item()

    df = preprpcessor.preprocess(df, user, item)

    y = df['rating']
    X = df.drop(['rating'], axis=1)

    X_train, X_valid, y_train, y_valid = train_test_split(X, y,
                                                    test_size=0.2,
                                                    random_state=1,
                                                    stratify=y)

    lgb_regressor = LGBMRegressor()
    lgb_regressor.train(X_train, X_valid, y_train, y_valid )

すると、こんな感じに勝手にFeature importanceとかが取れたりします。

LGBMの木の構造を保存する

Feature importanceより実際の木の構造が分かったほうが個人的にはいろいろわかることもあると思ってます。（そんなことない？） LGBMだと、dtreevizを使うとbinary classificationもしくはregressionなら木の構造をいい感じに画像を出力できたりします。

これはモデルの中身の理解に有用な情報だと思われ、これも自動でトラッキングしたくなったりします。そんなときにはこんな感じで一度dtreevizで可視化した木の構造を保存してあげて、MLflowのArtifactとして保存してあげたら実験の結果としてトラッキングしてくれます。

import lightgbm as lgb
import sklearn.preprocessing as sp
from dtreeviz.trees import *
import mlflow
from sklearn.model_selection import train_test_split
import dataloader, preprpcessor


class LGBMRegressor:
    def __init__(self):
        pass

    def train(self, X_train, X_valid, y_train, y_valid):
        mlflow.lightgbm.autolog()

        train_data = lgb.Dataset(X_train, label=y_train)
        validation_data = lgb.Dataset(X_valid, label=y_valid, reference=train_data)

        params = {
            'objective': 'regression',      
            'metric': "rmse",               
            'learning_rate': 0.05,
            'num_leaves': 21,               
            'min_data_in_leaf': 3,    
            'early_stopping':100,    
            'num_iteration': 1000           
            }

        self.model = lgb.train(params,
               train_set=train_data,
               valid_sets=validation_data,
               verbose_eval=50)
        
        viz = dtreeviz(self.model,
                    x_data = X_valid,
                    y_data = y_valid,
                    target_name = 'rating',
                    feature_names = X_train.columns.tolist(),
                    tree_index = 0)

        filename = 'lgb_tree.svg'
        viz.save(filename)
        mlflow.log_artifact(filename) # ←ここ
        
    def predict(self):
        pass

if __name__ == '__main__':
    df = dataloader.data()
    user = dataloader.user()
    item = dataloader.item()

    df = preprpcessor.preprocess(df, user, item)

    y = df['rating']
    X = df.drop(['rating'], axis=1)

    X_train, X_valid, y_train, y_valid = train_test_split(X, y,
                                                    test_size=0.2,
                                                    random_state=1,
                                                    stratify=y)

    lgb_regressor = LGBMRegressor()
    lgb_regressor.train(X_train, X_valid, y_train, y_valid )