目前我正在使用 API 提供下一个单词预测模型。该模型在使用烧瓶时成功运行,但在使用 gunicorn 进行部署时解开对象存在问题。Pickeled 对象依赖于类定义,我在需要的地方明确地提供类定义。
class LanguageModel(nn.Module):
def __init__(self, vocab_size, embedding_size, hidden_size, n_layers=1, dropout_p=0.5):
# Defining layers
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):
# self.embed.weight = nn.init.xavier_uniform(self.embed.weight)
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)
# importing word indexes
with open(w2i, "rb") as f1:
word2index = pickle.load(f1)
with open(i2w, "rb") as f2:
index2word = pickle.load(f2)
# loading model
model = torch.load(wordModel)
def getNextWords(words):
results = []
data = [words]
data = flatten([co.strip().split() + ['</s>'] for co in data])
x = prepare_sequence(data, word2index)
x = x.unsqueeze(1)
x = batchify(x, 1)
with torch.no_grad():
hidden = model.init_hidden(1)
for batch in getBatch(x, 1):
inputs, targets = batch
output, hidden = model(inputs, hidden)
prob = output.exp()
word_id = torch.multinomial(prob, num_samples=1).item()
# word_probs = torch.multinomial(prob, num_samples=1).probs()
word = index2word[word_id]
results.append(word)
return [res for res in results if res.isalpha()][:4] # return results
app = Flask(__name__)
@app.route('/')
def home():
return "Home"
@app.route('/getPredictions', methods=["POST"])
def getPredictions():
#...... code .........
resultJSON = {'inputPhrase': inputPhrase,
'predictions': predictions} # predictions [nextPhrase]
print('result: ', predictions)
return jsonify(resultJSON)
if __name__ == '__main__':
app.run(host='0.0.0.0', port=3001, debug=True) # 10.2.1.29
Gunicorn wsgi.py 文件:
from m_api import app
import torch
import torch.nn as nn
from torch.autograd import Variable
if __name__ == "__main__":
class LanguageModel(nn.Module):
def __init__(self, vocab_size, embedding_size, hidden_size, n_layers=1, dropout_p=0.5):
# Defining layers
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):
# self.embed.weight = nn.init.xavier_uniform(self.embed.weight)
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)
app.run()
这个应用程序在由烧瓶提供服务时运行得非常好,但是当我使用 gunicorn 时会抛出一个错误:
model = torch.load(wordModel)
File "/home/.conda/envs/sppy36/lib/python3.6/site-packages/torch/serialization.py", line 426, in load
return _load(f, map_location, pickle_module, **pickle_load_args)
File "/home/.conda/envs/sppy36/lib/python3.6/site-packages/torch/serialization.py", line 613, in _load
result = unpickler.load()
AttributeError: Can't get attribute 'LanguageModel' on <module '__main__' from '/home/.conda/envs/sppy36/bin/gunicorn'>
为了解决这个问题,我也在 wsgi.py 文件中包含了类定义,但是在加载 pickeled 文件时仍然无法获取类定义。我需要在哪里指定类定义仍然未知。