import torch
import torch.nn as nn


# 定义一个模型
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.embedded = nn.Embedding(4, 128)
        self.rnn = nn.GRU(128, 1024, batch_first=True)
        self.linear = nn.Linear(1024, 10)
        self.dropout = nn.Dropout(p=0.1)

    def forward(self, x):
        x, hn = self.rnn(self.embedded(x))
        return self.dropout(self.linear(x))


def test_model():
    # 创建模型实例
    model1 = Model()
    model1.to(device='cpu')
    print("model1-->", model1)
    # 初始化数据
    input = torch.tensor([[1, 3, 2, 0]])
    # 把数据送给模型（未量化）
    result1 = model1(input)
    # print("result1-->", result1)
    # 模型保存
    torch.save(model1.state_dict(), './model_without_quantize.pth')

    # 创建量化模型实例
    # model：原始模型
    # qconfig_spec：待量化的层参数
    # dtype：量化权重的目标类型
    model2 = torch.quantization.quantize_dynamic(model=model1,
                                                 qconfig_spec={torch.nn.Linear, nn.GRU},
                                                 dtype=torch.qint8)
    print("model2-->", model2)
    torch.save(model2.state_dict(), './model_with_quantize.pth')
    # 把数据送给模型（已量化）
    result2 = model2(input)
    # print("result2-->", result2)


if __name__ == '__main__':
    test_model()