Сохранение склеарна `FunctionTransformer` с функцией обертывания - PullRequest
Новая
       42

Сохранение склеарна `FunctionTransformer` с функцией обертывания

0 голосов
/ 02 января 2019

Я использую 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 Ответ

0 голосов
/ 05 января 2019

Мне удалось взломать решение с помощью модуля marshal (в дополнение к pickle) и переопределить магические методы getstate и setstate, используемые pickle. 

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")

Это работает, удалив функцию f из рассылки и вместо marshaling его код и его имя.

dill также выглядит как хорошая альтернатива маршалингу

...