グラフ機械学習と強化学習について

主にグラフ機械学習や強化学習手法を記載します。

ハイパーパラメータ管理（２）：MLflow + Optuna

前回の記事に引き続き、MLflow用いた機械学習のライフサイクルの管理をしていきます。

今回はグリッドサーチではなくベイズ最適化によるハイパーパラメータ探索を行っていきます。また、MLflowのモデルバージョン管理機能を用いていきます。

Optuna

Preferred Networkが開発しているハイパーパラメータチューニングのフレームワークです。私はRay Tuneを使うことが多いのですが、日本語によるサンプルコードが多いのと社内でよく使われることもあって試してみます。mlflowとも連携ができますが、各runごとにモデルのアーティファクトを保存したい場合、少し使いにくいです。

github.com

サンプルコード

LightGBMのハイパーパラメータを探索します。まずは、目的関数を設定します。

def objective(trial):
    # CVの方が良いが簡単のためホールドアウト
    data, target = sklearn.datasets.load_breast_cancer(return_X_y=True)
    train_x, test_x, train_y, test_y = train_test_split(data, target, test_size=0.25)
    dtrain = lgb.Dataset(train_x, label=train_y)

    # ここはhydraで管理した方が良いかも？
    # https://github.com/optuna/optuna/tree/master/examples/hydra
    param = {
        "objective": "binary",
        "metric": "binary_logloss",
        "verbosity": -1,
        "boosting_type": "gbdt",
        "lambda_l1": trial.suggest_loguniform("lambda_l1", 1e-8, 10.0),
        "lambda_l2": trial.suggest_loguniform("lambda_l2", 1e-8, 10.0),
        "num_leaves": trial.suggest_int("num_leaves", 2, 256),
        "feature_fraction": trial.suggest_uniform("feature_fraction", 0.4, 1.0),
        "bagging_fraction": trial.suggest_uniform("bagging_fraction", 0.4, 1.0),
        "bagging_freq": trial.suggest_int("bagging_freq", 1, 7),
        "min_child_samples": trial.suggest_int("min_child_samples", 5, 100),
    }

    gbm = lgb.train(param, dtrain)
    preds = gbm.predict(test_x)
    pred_labels = np.rint(preds)
    accuracy = sklearn.metrics.accuracy_score(test_y, pred_labels)
    return accuracy

次に、mlflowのコールバック関数を記載します。

def mlflow_callback(study, trial):
    trial_value = trial.value if trial.value is not None else float("nan"))
    with mlflow.start_run(experiment_id=experiment_id, run_name=study.study_name):
        mlflow.log_params(trial.params)
        mlflow.log_metric("accuracy", trial_value)

各トライアルにおいてハイパラとメトリックの結果を保存するように設定します。あとは実行するだけです。

if __name__ == "__main__":
    # 実験名と同じ階層の実験idを抽出する
    experiment_name = 'optuna'
    client = mlflow.tracking.MlflowClient(tracking_uri=mlflow.get_tracking_uri())
    for exp in client.list_experiments():
        if experiment_name == exp.name:
            experiment_id = exp.experiment_id
            break
    else:
        experiment_id = client.create_experiment(experiment_name)

    # ベイズ最適化の実行
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=100, callbacks=[mlflow_callback])

    print("Number of finished trials: {}".format(len(study.trials)))
    print("Best trial:")
    trial = study.best_trial
    print(f"Value: {trial.value}")
    print("Params:")
    for key, value in trial.params.items():
        print(f"{key}: {value}")

実行結果は以下のようになります。

[I 2021-03-08 22:56:28,200] A new study created in memory with name: no-name-84a20927-29be-4994-bd3b-cc311d81fd17
[I 2021-03-08 22:56:28,247] Trial 0 finished with value: 0.972027972027972 and parameters: {'lambda_l1': 4.215060077417707e-05, 'lambda_l2': 0.015415839449936278, 'num_leaves': 238, 'feature_frac
tion': 0.6673400621108638, 'bagging_fraction': 0.7442018650062918, 'bagging_freq': 2, 'min_child_samples': 41}. Best is trial 0 with value: 0.972027972027972.
[I 2021-03-08 22:56:28,487] Trial 1 finished with value: 0.958041958041958 and parameters: {'lambda_l1': 4.533990836656474, 'lambda_l2': 0.8259430522162956, 'num_leaves': 56, 'feature_fraction':
0.9744784080586081, 'bagging_fraction': 0.5414331193881199, 'bagging_freq': 3, 'min_child_samples': 9}. Best is trial 0 with value: 0.972027972027972.

MLflow ui

得られた実行結果は次のように保存されています。

f:id:udnp:20210308231617j:plain
mlflow ui optuna

モデルのバージョン管理

さて、得られたデータからモデルのバージョン管理をしていきます。

f:id:udnp:20210308231811j:plain
MLflow Model

ここから少し手順が面倒になってきます。というのもartifactとmodelのメタ情報は別々に保存していく必要があるからです。

MLflow Trackingを見るとArtifactを保存するパターンは4つほどあります

  1. MLflow on localhost
  2. MLflow on localhost with SQLite
  3. MLflow on localhost with Tracking Server
  4. MLflow with remote Tracking Server, backend and artifact stores

シナリオ1の場合は、特に何も考えずに実行したらMLflow側でmlrunsのフォルダを作成してくれます。experiment_idを指定してあげれば実験ごとにフォルダを作成してくれます。 ここでは、シナリオ２の、localhost + SQLiteの場合を試してみます。

MLflow on localhost with SQLite

ここで必要なのはSQLAlchemyです。データベースのスキーマはMLflow側で作成してくれるので、空のDBを作成します。

from sqlalchemy import create_engine

engine = create_engine('sqlite:///mlruns.db', echo=True)

次にMLflowのサンプルコードを実行しています。

def run():
    # Create two runs Log MLflow entities
    mlflow.set_tracking_uri('sqlite:///mlruns.db')
    with mlflow.start_run() as run1:
        params = {"n_estimators": 3, "random_state": 42}
        rfr = RandomForestRegressor(**params).fit([[0, 1]], [1])
        mlflow.log_params(params)
        mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model")

    with mlflow.start_run() as run2:
        params = {"n_estimators": 6, "random_state": 42}
        rfr = RandomForestRegressor(**params).fit([[0, 1]], [1])
        mlflow.log_params(params)
        mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model")

    # Register model name in the model registry
    name = "RandomForestRegression"
    client = mlflow.tracking.MlflowClient()
    client.create_registered_model(name)
　
　# 名前が重複する場合はエラーになるので例外処理を追加しておく
    try:
        client.get_registered_model(name)
    except:
        client.create_registered_model(name)

    # Create a two versions of the rfr model under the registered model name
    for run_id in [run1.info.run_id, run2.info.run_id]:
        model_uri = "runs:/{}/sklearn-model".format(run_id)
        mv = client.create_model_version(name, model_uri, run_id)
        print("model version {} created".format(mv.version))
    print("--")

    # Fetch the last version; this will be version 2
    mv = client.get_model_version(name, mv.version)
    print_model_version_info(mv)

これが行っているのは、各runごとにRandomForestを作成して、各パラメータと学習済みモデルを保存しています。それをmlflow.tracking.MlflowClient()を使って、モデルを登録していきます。ここで大事なのは、初めにmlflow.set_tracking_uri('sqlite:///mlruns.db')でtracking先をdbにしていることです。

f:id:udnp:20210308233523j:plain
ER図

ER図を見てもわかるようにartifact_uriがモデルを参照する場所になっています。今回ではローカルに保存していますが、AWS S3などにも簡単に保存できるようになっています。

面倒なのは、clientとmlflow.run_start()です。clientの場合はrun_idを作成できるのですが、フォルダが作成されないので、最初にwith mlflow.run_start():を実行してあげる必要があります。この操作を行わないとくMLflowのUI上からアーティファクトを参照することができなくないため、エラーとなっててしまいます。

optunaの場合はベストなハイパラを使って再度学習させる必要があるので、最後にモデルを保存するようなコードを追加すればokです。

if __name__ == '__main__':
    data = sklearn.datasets.load_breast_cancer(return_X_y=False)
    train_x, test_x, train_y, test_y = train_test_split(data.data, data.target, test_size=0.25)
    dtrain = lgb.Dataset(train_x, label=train_y)
    
    # optunaで得られたベストなパラメータを使って学習
    lgb_model = lgb.train(trial.params, dtrain)
    preds = lgb_model.predict(test_x)
    pred_labels = np.rint(preds)
    accuracy = sklearn.metrics.accuracy_score(test_y, pred_labels)

    # データのスキーマを登録する
    df = pd.DataFrame(data.data, columns=data.feature_names)
    signature = infer_signature(df, preds)

    with mlflow.start_run(experiment_id=experiment_id) as run:
        mlflow.log_params(trial.params)
        mlflow.log_metric('accuracy', accuracy)
        mlflow.lightgbm.log_model(lgb_model, artifact_path='lightgbm-model', signature=signature)

    name = 'LightGBM'
    tags = {'data': 'breast cancer'}
    try:
        client.get_registered_model(name)
    except:
        client.create_registered_model(name)

    run_id = run.info.run_id
    model_uri = "runs:/{}/lightgbm-model".format(run_id)
    mv = client.create_model_version(name, model_uri, run_id, tags=tags)
    print("model version {} created".format(mv.version))

このようにすることでモデルのバージョン管理を行うことができます。モデルの精度が悪化してきた場合は、このスクリプトを実行することで継続的にメンテナンスできます。

デプロイ

mlflowはモデルをデプロイする機能もあり、REST APIで予測結果を返すこともできます。仮想環境やDocker環境を自動で構築してくれます。

mlflow serve models -m <モデルのディレクトリ>

予測を行う場合はPOSTを行います。

import requests
url = 'http://127.0.0.1:5000/invocations'
headers = {'Content-Type': 'application/json'; 'format': 'pandas-split'}

r = requests.post(url, df.to_json('split'), headers)

もしくは、単純に以下の形式でもよいです。

curl -X POST -H "Content-Type:application/json; format=pandas-split" --data "{\"columns\":[\"alcohol\", \"chlorides\", \"citric acid\", \"density\", \"fixed acidity\", \"free sulfur dioxide\", \"pH\", \"residual sugar\", \"sulphates\", \"total sulfur dioxide\", \"volatile acidity\"],\"data\":[[12.8, 0.029, 0.48, 0.98, 6.2, 29, 3.33, 1.2, 0.39, 75, 0.66]]}" http://127.0.0.1:1234/invocations

モデルの読み込みも簡単です。

loaded_model = mlflow.pyfunc.load_model(mlflow_pyfunc_model_path)

各MLflowが提供しているModelAPIを使ってもよいと思います。

最後に

MLflowを使うとモデル管理を本格的に行うことができます。Microsoft Azure MLやAmazon SageMaker等にデプロイができるようなので、次回はSageMaker (or Kubernetes)でデプロイをやってみようと思います。

実際の業務でモデル選定やハイパラ調整で時間を大きく割くのは無駄（記述子・特徴量生成の方が大事）なので、簡単なテーマならさくっとモデルを構築してすぐプロダクトを提供できるような形にしていきたいです。