LightGBM + Optuna でベストモデル保存

さいきん書いたんで、どういうコードを書いたかをテキトーに書く。

ちなみに、公式のモデル保存の例はこちら。

https://optuna.readthedocs.io/en/stable/faq.html#how-to-save-machine-learning-models-trained-in-objective-functions

嫌だ、trial数分のモデルを保存するだなんて！

では、話を戻して。。。

今回の実装の制限事項は、シリアルでしか正しく動かないことかな。

準備するもの

Objectiveクラス
Callbackクラス

Objectiveの実装

Objectiveクラスは、__call__(self, trial)を実装。
__init__で、X_trainなどに加えて、Callbackクラスも受け取る。
__call__(self, trial)内で、学習したのち、良い悪いに関わらず、出来たモデルをCallbackクラスに引き渡す、その際のキーは、trialのnumberにする。
max_binもチューニングしたいなら、lgb.Datasetは毎回作らないといけない。

class Objective:
    def __init__(self, X_train, y_train, X_test, y_test, callback):
        self.X_train = X_train
        self.y_train = y_train
        self.X_test = X_test
        self.y_test = y_test
        self.callback = callback

    def __call__(self, trial):
        # do something
        params = self.create_params(trial)
        train_set = lgb.Dataset(self.X_train, self.y_train)
        valid_set = lgb.Dataset(self.X_test, self.y_test, reference=train_set)
        pruning_callback = optuma.integration.LightGBMPruningCallback(trial, 'binary_logloss')
        model = lgb.train(params, train_set, valid_sets=[valid_set], callbacks=[pruning_callback])
        self.callback.register_model(trial.number, model)
        y_test_pred = np.rint(model.predict(self.X_test))
        auc_test = roc_auc_score(self.y_test, y_test_pred)
        return 1 - auc_test
        
    def create_params(self, trial):
        return something trial suggestion code

Callbackの実装

Callbackクラスは元来、study.optimize()に渡すものです。今回も渡しますけどね。
__call__(self, study, trial)を実装。
study.best_trial.numberがtrial.numberと一致したら、ベストが更新されてます。だったら、登録されてるはずのモデルを取り出せるはずです。ベストが更新されてないなら、要らないモデルが登録されてるはずです、消してしまいましょう。
モデルの登録はただのdictを使うだけです。
ついでに、bestが更新されたら、best_paramsをjsonにして保存したり、更新されなくても、studyのdataframeを毎回保存したりなんかも出来ます。これをしてると、何トライアル目かを高速に流れるログから探さなくてすみます。

class Callback:
    def __init__(self):
        self.models = {}
    def register_model(self, trial_number, model):
        self.models[str(trial_number)] = model
    def unregister_model(self, trial_number):
        self.models.pop(str(trial_number), None)
    
    def unregister_other_model(self, trial_number):
        model = self.models.pop(str(trial_number), None)
        self.models.clear()
        self.models[str(trial_number)] = model
    def get_model(self, trial_number):
        return self.models[str(trial_number)]
    def __call__(self, study, trial):
        if study.best_trial.number == trial.number:
            self.unregister_other_model(study.best_trial.number)
            self.save(study)
        else:
            self.unregister_model(trial.number)
            # save study.trials_dataframe()
    def save(self, study):
        model = self.get_model(study.best_trial.number)
        # save model
        # save study.best_params
        # save study.trials_dataframe()

思ったより終わらずに、止めたくなることがある。そういう時に備えて、ベスト更新時に、ファイル保存。ベスト未更新でもtrials_dataframe()を毎回ファイル保存。

利用例

callback = Callback()
objective = Objective(X_train, y_train, X_test, y_test, callback)
sampler = optuna.samplers.TPESampler(seed=1)
study = optuna.study.create_study(sampler=sampler)
study.optimize(objective, n_trials=100, callbacks=[callback])

best_model = callback.get_model(study.best_trial.number)

再現性を持たせたい時には、TPESamplerのseedを固定する必要がある。

ベストモデルを毎回保存するなら、Callbackの__call__の中で書けばいい。
最後にcallbackから取り出して保存するなら、optimize後に取り出して保存すればいい。

add_modelを明示的に呼ぶのがやはりださい。

callback = Callback()
objective = Objective(X_train, y_train, X_test, y_test, lambda x, y: callback.add_model(x, y))

こっちのパターンのほうが汎用性が高いか？
結局はobjective側で何を登録するのか決めてるので、利点はなさそう。