FLAMLで学習済みのモデルを保存する・読み込む方法

FLAMLでAutoMLしてモデルを作成したあと、そのモデルを保存して読み込む方法が分からず、調べたのでメモ。以下の2種類がありそう。

1. Pickle等でAutoMLのオブジェクトとしてそのまま保存する
2. 実際に使用された学習モデルのパッケージの重みとして保存する

学習部分

from flaml import AutoML
from sklearn.model_selection import train_test_split
from sklearn.datasets import fetch_california_housing
from sklearn.metrics import mean_squared_error

automl = AutoML()
automl_settings = {
    "time_budget": 10,  # in seconds
    "metric": 'r2',
    "task": 'regression',
    "log_file_name": "california.log",
}
X, y = fetch_california_housing(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

automl.fit(X_train=X_train, y_train=y_train,**automl_settings)
print(automl.model.estimator)

Pickle等でAutoMLのオブジェクトとしてそのまま保存する

AutoMLのオブジェクトで読み出す場合にこちらを使用する。

import pickle
with open('automl.pkl', 'wb') as f:
    pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)

automl_pkl = pickle.load(open('automl.pkl', 'rb'))
y_pred = automl_pkl.predict(X_test)
print(mean_squared_error(y_pred, y_test, squared=False))

実際に使用された学習モデルのパッケージの重みとして保存する

推論時にはAutoMLを使用せずに、ベストな学習モデルのネイティブなパッケージを使用して推論したい場合にこちらを使用する。実際に何の学習モデルが選択されたかを、予め知る必要がある。

from catboost import CatBoostRegressor

automl.model.estimator.save_model("best_model.txt")

catboost = CatBoostRegressor()
catboost.load_model("best_model.txt")
y_pred = catboost.predict(X_test)
print(mean_squared_error(y_pred, y_test, squared=False))