序列化“腌制” sklearn模型时包括自定义类吗？ ...模块'__main__'没有属性

Question

AttributeError: module '__main__' has no attribute 'PandasDataFrameSelector'

如何解决此问题，以确保在腌制模型时包含自定义类？

我有一个自定义sklearn类，用于在管道中选择布尔cols，并希望通过api端点进行调用，但出现错误AttributeError: module '__main__' has no attribute 'PandasDataFrameSelector'。

在Jupyter笔记本中进行测试时，我没有问题，但是当我拥有.py时，api端点尝试读取.pkl时会出现错误，因为自定义类未包含在.pkl中）

我需要在腌制模型之前或之后添加if __name__ == '__main__':吗？

我是python dev ops部署的新手，只能找到本文=> https://www.stefaanlippens.net/python-pickling-and-dealing-with-attributeerror-module-object-has-no-attribute-thing.html

from sklearn.externals import joblib
from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np
import pandas as pd

# model training here + custom class ...

class PandasDataFrameSelector(BaseEstimator, TransformerMixin):
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        return X[self.attribute_names].values

# export model
joblib.dump(pl, "model.pkl")

# read model
class PandasDataFrameSelector(BaseEstimator, TransformerMixin):
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names
    def fit(self, X, y=None):
        return self
    def transform(self, X):
        return X[self.attribute_names].values

model = joblib.load("model.pkl")