使用Gunicorn部署Flask API加载pkl模型时遇AttributeError问题

这个问题的核心是pickle序列化模型时对函数命名空间的强依赖，我来帮你拆解问题并给出可行的解决办法：

问题根源 #

你训练这个出问题的模型时，map3函数是定义在**直接运行的脚本（也就是__main__模块）**里的。pickle保存模型时，会把这个函数的归属记录为__main__.map3。但用Gunicorn启动Flask应用时，进程的__main__模块变成了Gunicorn的入口文件（/Users/.../env/bin/gunicorn），这个模块里根本没有map3函数，所以加载模型时就会抛出找不到属性的错误。

其他模型能正常工作，说明它们依赖的自定义函数要么不在__main__模块里，要么没有依赖这类自定义函数。

解决方案 #

方案1：重构训练代码（推荐，一劳永逸）

把训练时用到的自定义函数（比如map3）移到一个单独的Python模块中，训练和加载时都从这个模块导入函数：

1. 创建一个工具模块，比如model_utils.py，把map3函数放在这里：

# model_utils.py
def map3(y_true, y_pred):
    # 这里写你原来的map3函数实现
    pass

2. 修改训练脚本，从这个模块导入map3再训练模型：

# train_script.py
from model_utils import map3
# 执行模型训练逻辑，使用map3函数
# 保存模型
import pickle
pickle.dump(model, open("problem_model.pkl", "wb"))

3. 在Flask的加载模型代码中，同样确保能导入model_utils.py里的map3（只要这个模块在应用的Python路径下即可），之后加载模型就不会有问题了。

方案2：不重新训练，在加载时修复命名空间

如果不想重新训练模型，可以在Flask的WSGI入口文件（比如wsgi.py）中，手动把map3函数挂载到__main__模块：

# wsgi.py
from flask import Flask

# 先定义或导入和训练时完全一致的map3函数
def map3(y_true, y_pred):
    # 必须和训练时的实现完全相同，不能有改动
    pass

# 把map3挂载到__main__模块，让pickle能找到对应属性
import __main__
__main__.map3 = map3

# 然后加载模型
import pickle
model = pickle.load(open("problem_model.pkl", "rb"))

app = Flask(__name__)
# 后续的API路由定义...

注意：这里的map3函数必须和训练时的实现完全一致，否则可能导致模型推理结果异常。

验证方法 #

你可以在训练时打印map3.__module__，确认它的归属模块：

print(map3.__module__)

如果输出是__main__，就说明会遇到这个问题；如果是model_utils这类自定义模块，就不会有问题。

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