pytorch 提供torch.quantization.quantize_dynamic函数,对模型中的某些层进行量化


1.初始化一个RNN模型,里面包含了LSTM层和全连接层,使用torch.quantization.quantize_dynamic对模型进行量化。

import torch.quantization

quantized_model = torch.quantization.quantize_dynamic(
    rnn, {nn.Linear}, dtype=torch.qint8                     #rnn为模型的名字,我们只量化线性层
)
print(quantized_model)

如果想量化线性层和LSTM层,将{nn.Linear}改为{nn.Linear,nn.LSTM}即可
dtype=torch.qint8 表示量化为有符号8位数,也可以选择无符号8位数quint8

2.检查量化后各层参数

在这里插入图片描述
以上是全连接层的参数,如图所示,全连接层的参数被量化为qint8(有符号8位数),zero_point和scale是pytorch量化时自己默认的量化参数,这两个参数的意思可以参考这篇博客torch.quantize_per_tensor()函数参数详解
在这里插入图片描述
如图,LSTM层未量化,所以LSTM层的参数是32位浮点数

3.什么是动态量化

torch.quantize_per_tensor()函数的scale和zero_point需要自己设定。
所谓动态是指这个函数torch.quantization.quantize_dynamic能自动选择最合适的scale和zero_point。
这两个参数的意思可以参考这篇博客torch.quantize_per_tensor()函数参数详解

4.量化后的训练和验证

量化后的模型,只能用于推理验证,不能训练!!!
执行梯度迭代操作时会报错。。。

5.总结

目前只想到了渐进量化的方式去减小量化误差。。。。。。。。。

等我想到更好的,可训练的量化方法再更新。。。。。。。。

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐