[FIXED] Unable to unpickle an object while predicting using gunicorn

Issue

Currently I am serving a next word prediction model using an API. The model was successfully working when using flask but there is an issue in unpickeling the object when using gunicorn for deployment. Pickeled object is dependent on class definition and I am supplying the class definition explicitly wherever it's needed.

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 file:

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()

This app when served by flask runs perfectly fine but when I use gunicorn an error is thrown out:

    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'>

To resolve this I included class definition in wsgi.py file too, but still it's not able to get the class definition at the time of loading the pickeled file. Where do I need to specify the class definition is still unknown.

Solution

The issue is because the gunicorn looks for the Class definition int's main method, i.e. the gunicorn executable file. That's why even the explicit definition of class in both the .py files did not do the expected work while running on gunicorn but did while using flask. To overcome this issue I explicitly defined the class in the gunicorn executable file and it worked. For now, I found this as the workable solution.

gunicorn.py

#!/home/user/anaconda3/envs/envName/bin/python

import re
import sys

from gunicorn.app.wsgiapp import run

import torch
import torch.nn as nn
from torch.autograd import Variable

USE_CUDA = torch.cuda.is_available()

if __name__ == '__main__':
    # defining model class
    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)

        def init_hidden(self, batch_size):
            hidden = Variable(torch.zeros(self.n_layers, batch_size, self.hidden_size))
            context = Variable(torch.zeros(self.n_layers, batch_size, self.hidden_size))
            return (hidden.cuda(), context.cuda()) if USE_CUDA else (hidden, context)

        def detach_hidden(self, hiddens):
            return tuple([hidden.detach() for hidden in hiddens])

        def forward(self, inputs, hidden, is_training=False):
            embeds = self.embed(inputs)
            if is_training:
                embeds = self.dropout(embeds)
            out, hidden = self.rnn(embeds, hidden)
            return self.linear(out.contiguous().view(out.size(0) * out.size(1), -1)), hidden


    sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
    sys.exit(run())

Answered By - Mousam Singh