XGBoostはpickleでなくJSONで保存せよ 備忘録

概要

GPU付きPCでモデルのTrainingを行いpredictはGPUなし(つまりCPUのみ)のPCで行う運用を想定する。
そのときのXGBoostモデルの保存方法について備忘録を残す。

• 実施時期: 2022年9月
• Python: conda 3.8.13
• XGBoost: py-xgboost, py-xgboost-gpu 1.6.1

問題

XGBoostはLightGBMに比べTrainingが遅いことは有名だ。GPUを使えばCPUよりも体感で4~5倍高速になり十分な速さだが、どのPCにもGPUが乗っているわけではない。
しかしCPUであってもPredictはLightGBMより速いのでdeployを考えるとどうしてもXGBoostを使いたくなる。老舗なため情報が豊富なこともありがたい。
ただモデル保存時、いつもの便利なpickleで保存してしまうと運用時に困ったことになってしまうので対策を残す。

Training at PC with GPU

XGBoostをGPUでTrainingするときは下記のようなコードとなる。
ここのポイントはモデルをJSONで保存する点であり、pickleで保存しないこと。

    param_wk['tree_method'] = 'gpu_hist'
    param_wk['predictor'] = 'cpu_predictor'    # ⇐　deploy先のPCにGPUがない運用
    reg = xgb.XGBRegressor(**param_wk)
    reg.fit(X_train, y_train, eval_set = [(X_train, y_train), (X_val, y_val)], hogehoge....)
    reg.save_model('reg_model.json')

Predict at PC without GPU

GPUありPCで保存したJSONファイルをGPUがないPCで読込んでpredictさせるときは下記となる。

    reg = xgb.XGBRegressor()
    reg.load_model('reg_model.json')
    reg.get_booster().set_param({'tree_method': 'hist'})    # ⇐　jsonファイルをエディタで編集してもOK
    y_pred = reg.predict(X_test)

JSONで保存したときは、それを読込んだあとでset_param関数で‘tree_method’をCPUを示す'hist'に変更することができる。変更しなければモデルはGPUを期待したままなので後述のwarningが発生する。
pickleで読み込んだときはset_param関数が使えないので詰む。
JSONファイルなのでエディタで開いてファイル冒頭に書かれている‘tree_method’: ‘gpu_hist’　を　‘tree_method’: ‘hist’　に書き換えてもよい。pickleはバイナリなのでそんなことはもちろんできない。
ちなみにpickleで保存してしまい、それを読み込むと下記のようなwarningが発生しモデルをロードすることができない。

WARNING: D:\bld\xgboost-split_1660208973102\work\src\gbm\gbtree.cc:390: Loading from a raw memory buffer on CPU only machine.  Changing tree_method to hist.

"Unknown tree updater grow_gpu_hist"と出るときもあるが、同様にロードはできない。

まとめ

Trainingもpredictも同じPCならもちろんpickleでもOKだが、JSONにするとモデルのハイパラやパラメータをエディタで確認できるので、基本的にはJSONで保存すべきだろうと、思った。
またpickleではXGBoostのバージョン互換性がなくなることも考えられ、やはりJSONの一択と思われる。

以上