如何在Flask应用中加载含自定义损失函数的Learner模型

问题原因 #

Pickle加载包含自定义函数的模型时，要求该函数必须存在于当前运行环境的命名空间中。你用flask.exe启动应用时，程序的主模块是Flask自身的__main__.py，而非你的Flask脚本，导致脚本里定义的combined_loss没被加载到这个命名空间，所以报错。

方案1：将自定义损失函数独立成模块（推荐） #

把combined_loss放到单独的Python文件里，训练和Flask应用都从这个模块导入，确保两边的函数完全一致：

1. 创建loss_utils.py文件，写入损失函数：

def combined_loss(preds, targets):
    # 这里是你原来的损失函数逻辑
    ...

2. 训练模型时，从该模块导入后再训练保存：

from loss_utils import combined_loss
# 训练并保存Learner模型的代码

3. Flask应用中，先导入函数再加载模型：

from flask import Flask
from loss_utils import combined_loss
import pickle

app = Flask(__name__)

# 在加载模型前确保函数已导入
with open('your_model.pkl', 'rb') as f:
    model = pickle.load(f)

# 后续的Flask路由逻辑

方案2：在Flask脚本加载模型前重定义损失函数 #

确保重定义的combined_loss和训练时的代码完全一致（包括参数、计算逻辑），再执行模型加载：

from flask import Flask
import pickle

app = Flask(__name__)

# 先重定义和训练时一模一样的combined_loss
def combined_loss(preds, targets):
    # 复制训练时的损失函数代码
    ...

# 现在加载模型
with open('your_model.pkl', 'rb') as f:
    model = pickle.load(f)

注意：如果后续修改损失函数，需要同时更新训练和Flask脚本里的代码，容易出现不一致，仅适合临时测试。

方案3：直接运行Flask脚本而非用flask命令 #

如果你的Flask脚本是app.py，不要用flask run启动，而是直接执行：

python app.py

这样程序的主模块就是你的app.py，脚本里定义的combined_loss会存在于__main__命名空间中，加载模型时就能找到。

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