Unable to unpickle an object while predicting using gunicorn

Question

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.

Answer 1

It is too late, but I hope someone finds this useful.

The reason why it does not work is because when you unpickle your own models, pickle stores the module and the name of the object. From what I see in your output, it is likely that you trained the model in a script where you define your own model in the same script, and therefore when you pickle your model it says that the module is __main__. The problem with that, is that when you run gunicorn, as you can expect __main__ is the script that runs gunicorn, and when you are unpiclking your object it looks for it in the gunicorn script (because the pickled file says it is there).

To overcome this, when you train your model, define a directory where your model is going to be, something like

directory/__init__.py directory/model.py

in __init__.py make sure to have something like

"from .model import ModelClass"

Now try to train the model importing it as in "from directory import ModelClass"

train it and the pickle it.

Then in gunicorn, in the .py that you unpickle the model you need to be sure that the directory/*py structure is kept, because when you do pickle.load it will try to look for something like

"directory.model .... ModelClass". You can actually print the first lines of the pickle file reading it as a binary file and you will see the kind of structure is has.

Advice: the best option and probably the best practice is to not pickle complex class objects, better stick to pickle objects that are installed via pip because they are in the default packages of python and when you pickle those you don't need to look out for that directory structure. So whenever possible, try to pickle objects from those classes, like sklearn models, dictionaries, lists, and then to load your model unpickle all those files and send them as parameters to the constructor.

Answer 2

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