Exporting AWD_LSTM to PyTorch for serving

geooff · February 6, 2021, 5:03pm

Hello all,

I’m wondering if anyone has any experience exporting trained AWD_LSTM models from FastAI to PyTorch for serving. I’m working in a situation where I need as small a serving image as possible so I’m trying to trim down my dependencies.

So far I’ve been prowling the forums and this is what I have.

learn = load_learner("path")
torch.save(learn.model.state_dict(), 'path')

config = {"my model params"}

model = torch.nn.LSTM(**config)
model.load_state_dict(torch.load('path'))
model.eval()

This gives me a bunch of key errors because my LSTM and AWD_LSTM are obviously different creatures.

Let me know if anyone has any experience with this. I’ve trained some great models with FastAI but now I’m struggling to make it to production given my circumstances.

ilovescience · February 7, 2021, 3:38am

Yeah you need to use the same AWD_LSTM model in fastai. As you can see, it is defined here:

github.com

fastai/fastai/blob/29be53d7ccaf7405320e2e47db8f35182abeff0a/fastai/text/models/awdlstm.py#L81


        if self.training and self.embed_p != 0:
            size = (self.emb.weight.size(0),1)
            mask = dropout_mask(self.emb.weight.data, size, self.embed_p)
            masked_embed = self.emb.weight * mask
        else: masked_embed = self.emb.weight
        if scale: masked_embed.mul_(scale)
        return F.embedding(words, masked_embed, ifnone(self.emb.padding_idx, -1), self.emb.max_norm,
                           self.emb.norm_type, self.emb.scale_grad_by_freq, self.emb.sparse)
# Cell
class AWD_LSTM(Module):
    "AWD-LSTM inspired by https://arxiv.org/abs/1708.02182"
    initrange=0.1
    def __init__(self, vocab_sz, emb_sz, n_hid, n_layers, pad_token=1, hidden_p=0.2, input_p=0.6, embed_p=0.1,
                 weight_p=0.5, bidir=False):
        store_attr('emb_sz,n_hid,n_layers,pad_token')
        self.bs = 1
        self.n_dir = 2 if bidir else 1
        self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token)
        self.encoder_dp = EmbeddingDropout(self.encoder, embed_p)

So in your code you will also have to include the same definitions and you then can load your model weights.

geooff · February 7, 2021, 3:59am

Thanks for your response!

When going through the source code I have a hard time determining where the PyTorch ends and the FastAI begins. Does the AWD_LSTM class itself yield a valid PyTorch model? Is it more or less just as easy as:

# Disclaimer: This is pseudo code and won't work
model = AWD_LSTM(**model_params)
model.load_state_dict(torch.load('test_model.pth'))
model.eval()

ilovescience · February 7, 2021, 4:48am

I am not sure myself, there might be some fastai-specific functions involved in the implementation but it’s unlikely. I would test it by copying and pasting the code into your script, running in a PyTorch-only environment, and see which functions it complains about as missing…