使用pickle保存函数返回的sktime包装类模型报错如何解决

报错原因 #

Python的pickle模块默认序列化自定义类时，依赖类在所属模块中的全局可引用路径：即类的__module__属性对应的模块，加上__qualname__属性对应的类名，要求反序列化时该路径下能找到对应的类定义。

你将SKtimeWrapper定义在sktime_wrapper函数内部时，该类属于函数的局部对象，它的__qualname__为sktime_wrapper.<locals>.SKtimeWrapper，并不在模块的顶层命名空间中，pickle序列化时找不到它的全局引用路径，就会抛出无法序列化局部对象的错误。而你把类定义放到模块顶层后，它的路径是全局可访问的，pickle就能正常处理。

解决方案 #

以下3种方案都能在保留通用包装能力的前提下解决问题：

方案1：动态注册内部类到全局命名空间 #

手动给动态生成的包装类设置正确的模块属性，并注册到当前模块的全局命名空间，让pickle能识别：

import pickle
import sys
from sktime.transformations.panel.rocket import Rocket
from sktime.datatypes._panel._convert import from_2d_array_to_nested

def sktime_wrapper(method_class):
    # 给每个包装类设置唯一类名，避免冲突
    class_name = f"SKtimeWrapper_{method_class.__name__}"
    # 动态生成继承类
    SKtimeWrapper = type(class_name, (method_class,), {})
    # 重写方法
    def transform(self, X):
        X = from_2d_array_to_nested(X)
        return super().transform(X)
    def fit(self, X, Y):
        X = from_2d_array_to_nested(X)
        return super().fit(X, Y)
    SKtimeWrapper.transform = transform
    SKtimeWrapper.fit = fit
    # 设置模块属性，让pickle认为它是顶层定义的类
    SKtimeWrapper.__module__ = __name__
    # 注册到当前模块的全局命名空间
    setattr(sys.modules[__name__], class_name, SKtimeWrapper)
    return SKtimeWrapper


model = sktime_wrapper(Rocket)
with open('model.pkl','wb') as f:
    pickle.dump(model, f)

该方案无需改变原有使用逻辑，兼容原生pickle的序列化/反序列化。

方案2：改用支持局部对象序列化的dill库 #

dill是pickle的扩展库，支持序列化局部类、闭包等原生pickle无法处理的对象，改动最小：

1. 安装依赖：pip install dill
2. 代码中替换导入即可，其余逻辑完全不变：

import dill as pickle
from sktime.transformations.panel.rocket import Rocket
from sktime.datatypes._panel._convert import from_2d_array_to_nested

def sktime_wrapper(method_class):
    class SKtimeWrapper(method_class):
        def transform(self, X):
            X = from_2d_array_to_nested(X)
            return super().transform(X)
        def fit(self, X, Y):
            X = from_2d_array_to_nested(X)
            return super().fit(X, Y)
    return SKtimeWrapper

model = sktime_wrapper(Rocket)
with open('model.pkl','wb') as f:
    pickle.dump(model, f)

注意反序列化时也需要用dill加载模型。

方案3：用组合模式代替动态继承 #

顶层定义通用包装类，把原sktime实例作为内部属性，不需要动态生成类：

import pickle
from sktime.transformations.panel.rocket import Rocket
from sktime.datatypes._panel._convert import from_2d_array_to_nested

class SKtimeWrapper:
    def __init__(self, method_class, *args, **kwargs):
        self.inner_model = method_class(*args, **kwargs)
    
    def transform(self, X):
        X = from_2d_array_to_nested(X)
        return self.inner_model.transform(X)
    
    def fit(self, X, Y):
        X = from_2d_array_to_nested(X)
        self.inner_model.fit(X, Y)
        return self
    
    # 可选：自动透传原类的其他属性和方法
    def __getattr__(self, item):
        return getattr(self.inner_model, item)

# 使用方式
model = SKtimeWrapper(Rocket)
with open('model.pkl','wb') as f:
    pickle.dump(model, f)

该方案逻辑清晰，序列化稳定性最高，适合需要长期维护的场景。

内容的提问来源于stack exchange，提问作者user6274417