基于pytorch的模型定点化

%matplotlib inline

(experimental) Dynamic Quantization on an LSTM Word Language Model

Author: James Reed <https://github.com/jamesr66a>_

Edited by: Seth Weidman <https://github.com/SethHWeidman/>_

Introduction

Quantization involves converting the weights and activations of your model from float
to int, which can result in smaller model size and faster inference with only a small
hit to accuracy.

In this tutorial, we’ll apply the easiest form of quantization -
dynamic quantization <https://pytorch.org/docs/stable/quantization.html#torch.quantization.quantize_dynamic>_ -
to an LSTM-based next word-prediction model, closely following the
word language model <https://github.com/pytorch/examples/tree/master/word_language_model>_
from the PyTorch examples.

# imports
import os
from io import open
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

1. Define the model

---

Here we define the LSTM model architecture, following the
model <https://github.com/pytorch/examples/blob/master/word_language_model/model.py>_
from the word language model example.

class LSTMModel(nn.Module):
    """Container module with an encoder, a recurrent module, and a decoder."""

    def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(LSTMModel, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(ntoken, ninp)
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        self.decoder = nn.Linear(nhid, ntoken)

        self.init_weights()

        self.nhid = nhid
        self.nlayers = nlayers

    def init_weights(self):
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, input, hidden):
        emb = self.drop(self.encoder(input))
        output, hidden = self.rnn(emb, hidden)
        output = self.drop(output)
        decoded = self.decoder(output)
        return decoded, hidden

    def init_hidden(self, bsz):
        weight = next(self.parameters())
        return (weight.new_zeros(self