加载含自定义预处理类的Pickle模型报错：无法找到PriceImputer

问题场景 #

• 项目结构：
  • notebooks/：用于训练模型（训练过程中使用了自定义预处理类PriceImputer）
  • src/main.py：尝试用joblib加载训练好的模型，已正确导入PriceImputer类，但运行时出现报错：

    AttributeError: Can't get attribute 'PriceImputer' on <module 'mp_main' from '...'>

原因 #

在notebooks中训练并保存模型时，joblib会记录自定义类的模块归属。如果直接在main.py的全局作用域执行模型加载逻辑，Python的多进程模块上下文（__mp_main__）无法匹配训练时的类模块路径，导致反序列化时找不到PriceImputer的定义。

解决办法 #

把模型加载的代码放到if __name__ == '__main__':代码块内：

import joblib
from your_preprocessing_module import PriceImputer  # 替换为实际导入路径

if __name__ == '__main__':
    # 模型加载逻辑
    trained_model = joblib.load('../notebooks/trained_model.pkl')
    # 后续模型调用代码

当脚本作为主程序运行时，__name__会被赋值为'__main__'，此时加载模型的上下文与脚本的模块上下文一致，joblib就能正确识别并找到PriceImputer类，完成模型反序列化。

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