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.总结
目前只想到了渐进量化的方式去减小量化误差。。。。。。。。。
等我想到更好的,可训练的量化方法再更新。。。。。。。。