問題

機械学習で特徴量を色々試行錯誤しながらモデルを学習してると、学習時に使った特徴量がわからなくなる。gitのコミットハッシュとか記録しておけば、特徴量はわかるけど。本番でロードするときに、そのバージョンのコードを使わないといけないから、合わせるのが面倒。複数バージョンを同時に使うのはもっと面倒

ソリューション

cloudpickleを使う https://github.com/cloudpipe/cloudpickle

普通のpickleが関数を参照(関数名とか？)で保存するのに対して、cloudpickleは一部の関数(自動識別 + 手動で指定)を値(ソースコード？)で保存できる。

以下のようなモジュールがあったとする。PickleTestClassの中身やfunc2の中身は試行錯誤で頻繁に書き換える。register_pickle_by_valueを使うと、強制的に値で保存させられる。

# pickle_test_class.py

import cloudpickle
import sys

# 値で保存させる
cloudpickle.register_pickle_by_value(sys.modules[__name__])

class PickleTestClass:
   def __init__(self):
       pass

   def func(self, x):
       return 4 * x

# numbaも行けた
@numba.njit
def func2(x):
   return 3 * x

学習でモデルを保存したとする。

import cloudpickle
from .pickle_test_class import PickleTestClass
pt = PickleTestClass()
data = cloudpickle.dumps(pt)
with open('/tmp/data.txt', 'wb') as f:
   f.write(data)

本番環境でモデルをロード。PickleTestClassは学習時から書き換えられているとする

import cloudpickle
with open('/tmp/data.txt', 'rb') as f:
   data = f.read()
pt = cloudpickle.loads(data)
print(pt.func(1))

↑、学習時のバージョンのPickleTestClassが動く

良くいじるコード(自分で書いたコード)は値で保存(学習時のコードが使われる)。いじらないコード(sklearnとかtalibとか)は参照(ロード時のコードが使われる)で保存。とするのが良さそう。

cloudpickle.register_pickle_by_value(自分で書いてる大元のモジュール)とすれば、その配下が全て対象になるらしい https://github.com/cloudpipe/cloudpickle/blob/6099fdb1229adca6b8a013896bc8ec8daabe73b0/cloudpickle/cloudpickle.py#L181

mlflow model

試してないけど、mlflowはcloudpickle使うっぽい。

https://github.com/mlflow/mlflow/blob/9c934f5a57cc8701475eb7ed0b6b072600b66bed/mlflow/sklearn/__init__.py#L432 で保存しているらしい。そのまま呼んでるだけ。log_modelでも自前でlog_artifactでもどっちでも良い気がする。

log_modelだと圧縮されないっぽい。 https://github.com/mlflow/mlflow/issues/2545