使用Gunicorn部署API时无法反序列化PyTorch模型对象的问题

我完全明白你踩的这个坑——用Flask跑模型好好的，换成Gunicorn就炸了，还报找不到LanguageModel类的错。这问题本质上是pickle序列化时记录的类路径和Gunicorn运行时的模块上下文不匹配导致的，咱们一步步来解决：

问题根源 #

当你用torch.save保存模型时，pickle会把类的完整模块路径（比如你原代码里的__main__.LanguageModel）存下来。但Gunicorn启动时，wsgi.py是作为模块被导入的（不是以__main__身份运行），而且你把LanguageModel定义在了if __name__ == "__main__"块里——这个块在Gunicorn加载wsgi.py时根本不会执行！就算你把类移到块外面，模型里记录的类路径是原代码的__main__.LanguageModel，而Gunicorn里类的路径是wsgi.LanguageModel，还是对不上，自然找不到。

最优解决方案：统一类的定义位置 #

把LanguageModel类单独放到一个独立的模块里，让所有需要它的地方（训练代码、Flask服务、Gunicorn的wsgi入口）都从这个模块导入，这样类的路径就完全一致了。

步骤1：创建单独的模型定义文件model_defs.py #

把LanguageModel类移到这个文件里，内容如下：

import torch
import torch.nn as nn

class LanguageModel(nn.Module):
    def __init__(self, vocab_size, embedding_size, hidden_size, n_layers=1, dropout_p=0.5):
        super(LanguageModel, self).__init__()
        self.n_layers = n_layers
        self.hidden_size = hidden_size
        self.embed = nn.Embedding(vocab_size, embedding_size)
        self.rnn = nn.LSTM(embedding_size, hidden_size, n_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(dropout_p)
    
    def init_weight(self):
        # 注意：这里的new_w需要确保在这个模块里能访问到，或者改成参数传入
        self.embed.weight.data.copy_(torch.from_numpy(new_w))
        self.linear.weight = nn.init.xavier_uniform_(self.linear.weight)
        self.linear.bias.data.fill_(0)

    # 补上你原代码调用的init_hidden方法，否则加载后调用会报错
    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data
        return (torch.zeros(self.n_layers, batch_size, self.hidden_size),
                torch.zeros(self.n_layers, batch_size, self.hidden_size))

步骤2：修改原训练/模型保存代码（可选，若已训练好可跳过） #

训练时从model_defs.py导入LanguageModel，再保存模型，确保模型记录的类路径是统一的：

from model_defs import LanguageModel

# ... 你的训练代码逻辑 ...

# 保存模型
torch.save(model, wordModel)

步骤3：修改Flask服务文件m_api.py #

从model_defs.py导入LanguageModel，再加载模型：

from flask import Flask, jsonify, request
import pickle
import torch
from model_defs import LanguageModel  # 关键：从统一模块导入类
# 导入你需要的其他工具函数（flatten、prepare_sequence等）

# 加载词表
with open(w2i, "rb") as f1:
    word2index = pickle.load(f1)
with open(i2w, "rb") as f2:
    index2word = pickle.load(f2)

# 加载模型（现在能正确找到LanguageModel类了）
model = torch.load(wordModel)
model.eval()  # 别忘了设置为评估模式

def getNextWords(words):
    # ... 你的原代码逻辑不变 ...

app = Flask(__name__)
@app.route('/')
def home():
    return "Home"
@app.route('/getPredictions', methods=["POST"])
def getPredictions():
    inputPhrase = request.json.get('inputPhrase')  # 补充获取输入的代码
    predictions = getNextWords(inputPhrase)
    resultJSON = {'inputPhrase': inputPhrase, 'predictions': predictions}
    print('result: ', predictions)
    return jsonify(resultJSON)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=3001, debug=True)

步骤4：简化wsgi.py文件 #

现在wsgi.py不需要再定义LanguageModel了，直接导入app即可：

from m_api import app

if __name__ == "__main__":
    app.run()

验证运行 #

现在用Gunicorn启动：

gunicorn --workers=4 --bind=0.0.0.0:3001 wsgi:app

应该就能正常加载模型，不会再报找不到LanguageModel的错误了。

为什么原来的方法不行？ #

1. 你把LanguageModel放在wsgi.py的if __name__ == "__main__"块里，Gunicorn加载wsgi.py时是作为模块导入，__name__不是__main__，所以这个类根本没被定义。
2. 就算移到块外面，模型保存时记录的类路径是原代码的__main__.LanguageModel，而Gunicorn里类的路径是wsgi.LanguageModel，pickle找不到对应的类，所以报错。

把类统一放到单独的模块里，所有地方都从这个模块导入，就能保证类的路径完全一致，pickle加载时自然能找到。

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