用它包装的函数保存一个 sklearn `FunctionTransformer`

Posted 2023-02-16

【中文标题】用它包装的函数保存一个 sklearn `FunctionTransformer`

【英文标题】：Saving an sklearn `FunctionTransformer` with the function it wraps

【发布时间】：2019-05-29 11:32:10

【问题描述】：

我正在使用 sklearn 的 Pipeline 和 FunctionTransformer 和自定义函数

from sklearn.externals import joblib
from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline

这是我的代码：

def f(x):
    return x*2
pipe = Pipeline([("times_2", FunctionTransformer(f))])
joblib.dump(pipe, "pipe.joblib")
del pipe
del f
pipe = joblib.load("pipe.joblib") # Causes an exception

我得到这个错误：

AttributeError: 模块 '__ main__' 没有属性 'f'

如何解决？

请注意，pickle 中也会出现此问题

【参考方案1】：

我能够使用marshal 模块（除了pickle）破解解决方案并覆盖pickle 使用的魔术方法getstate 和setstate。

import marshal
from types import FunctionType
from sklearn.base import BaseEstimator, TransformerMixin

class MyFunctionTransformer(BaseEstimator, TransformerMixin):
    def __init__(self, f):
        self.func = f
    def __call__(self, X):
        return self.func(X)
    def __getstate__(self):
        self.func_name = self.func.__name__
        self.func_code = marshal.dumps(self.func.__code__)
        del self.func
        return self.__dict__
    def __setstate__(self, d):
        d["func"] = FunctionType(marshal.loads(d["func_code"]), globals(), d["func_name"])
        del d["func_name"]
        del d["func_code"]
        self.__dict__ = d
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        return self.func(X)

现在，如果我们使用MyFunctionTransformer 而不是FunctionTransformer，代码将按预期工作：

from sklearn.externals import joblib
from sklearn.pipeline import Pipeline

@MyFunctionTransformer
def my_transform(x):
    return x*2
pipe = Pipeline([("times_2", my_transform)])
joblib.dump(pipe, "pipe.joblib")
del pipe
del my_transform
pipe = joblib.load("pipe.joblib")

其工作方式是从 pickle 中删除函数 f，取而代之的是 marshaling 其代码和名称。

dill 看起来也是一个很好的编组替代方案

【讨论】：

应该是：del my_transform 而不是 del f。这仍然适用于多个自定义函数或嵌套管道吗？

真的，谢谢，我修复了代码 sn-p。它适用于嵌套管道和任何可编组的东西（不是每个函数都是）​​

您确实打算在单独的脚本中加载您的管道，不是吗？因此，即使使用您当前的方法，您是否不需要在内存中的某处准备好MyFunctionTransformer 的代码或在调用joblib.load 之前导入？这比在导入中准备好function f 的代码更好。也许来自另一个脚本？我错过了什么吗？

您是否同意如果FunctionTransformer 将通过我的添加（即setstate 和getstate）实现，那么酸洗将包括所有管道所需的依赖项?

关于sklearn，当您使用pickle 和TfidfVectotizer 转换器时，您希望它存储vocab、tf 和idf 以便工作。我认为FunctionTransformer，它的唯一目的是用transformer 包装一个函数，至少应该保存这个函数，或者在不可能的情况下发出警告。附：根据这次讨论，我编辑了我的问答。