Pytorch 1.10.2 下模型量化踩坑

---

前言

本文主要记录自己在使用Pytorch==1.10.2下对模型进行量化的过程，以及量化过程中遇到的坑，如果想详细了解，可以去看知乎大佬的详细介绍。

首先介绍本文使用的大概模型结构，如下所示是一个3D CNN头部+4层2D CNN组成的前端模型，然后会输入到GRU中进行序列建模，最后通过一个全连接层获得输出结果。

下面会根据这种网络结果，分别对训练完成后的模型使用静态量化和动态量化。

class model:
	def __init__(self):
		self.front3D = nn.sequential(
								nn.Conv3d(),
								nn.BatchNorm3d(),
								nn.ReLU(),
								nn.MaxPool3d()
								)
		self.front = ResNet18()
		self.back = nn.GRU()
		self.fc = nn.Linear()
	def forward(self, x):
		...
		x = self.front3D(x)
		...
		x = self.front(x)
		x, _ = self.back(x)
		x = self.fc(x)
		return x

---

一、动态量化

• 动态量化那是真滴简单

1. 动态量化针对哪些？

	- RNN
	- LSTM
	- GRU
	- Linear

2.如何使用？

以实例代码为例

# 首先定义模型
model = Model()
# 使用动态量化接口进行量化，这里表示会将GRU和Linear量化
model = torch.quantization.quantize_dynamic(
    model, {nn.GRU, nn.Linear}, dtype=torch.qint8)

到这里就完了，然后可以直接进行推理，打印量化后的模型如下

...
(fc): DynamicQuantizedGRU(..., dtype=torch.qint8, ...)
(fc): DynamicQuantizedLinear(..., dtype=torch.qint8, ...)
...

二、静态量化

• 静态量化那是真滴麻烦

1. 静态量化针对哪些？

	- Conv2d, BatchNorm2d
	- Conv3d, BatchNorm3d
	- ReLU
	- Linear
	- ...
	基本上大部分CNN的算子都支持了，但是有例外，后面会提到。

2.如何使用？

以实例代码为例

# 首先定义模型
model = Model()
# 设置qconfig
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
# 融合网络层
listmix =  [
     ['front3D.0', 'front3D.1', 'front3D.2'],
     ['front.layer1.0.conv1', 'front.layer1.0.bn1'],
     ['front.layer1.0.conv2', 'front.layer1.0.bn2', 'front.layer1.0.relu'],
     ...
     ]

在这里，对于sequential封装的conv3d, bn3d和relu，直接使用索引即可，网络融合必须以下例子中：
- Convolution, Batch normalization
- Convolution, Batch normalization, Relu
- Convolution, Relu
- Linear, Relu
- Batch normalization, Relu

需要注意的是，如果是单独的ReLU和单独的BN是不能融合，此外单独的BN运行时会报错，后面会提。

# 融合
torch.quantization.fuse_modules(model, listmix, inplace=True)
# 使用实际样本，让量化模型观察数据分布
torch.quantization.prepare(model, inplace=True)
data = torch.randn(1, 44, 44)	# 读入的实际数据，这里是举个例子
data = data.unsequence(0)	# 需要增加一维
model(data)
torch.quantization.convert(model, inplace=True)

这个时候模型就转换好了，但是还需要再模型内添加量化和解量化的模块，分别是QuantStub()和DeQuantStub()，通过QuantStub()后，会从CPU后端变成QuantizeCPU，这样才能跑量化后的一些算子。代码修改如下：

from torch.quantization import DeQuantStub, QuantStub
class model:
	def __init__(self):
		self.front3D = nn.sequential(
								nn.Conv3d(),
								nn.BatchNorm3d(),
								nn.ReLU(),
								nn.MaxPool3d()
								)
		self.front = ResNet18()
		self.back = nn.GRU()
		self.fc = nn.Linear()
		self.quant = QuantStub()
		self.dequant = DeQuantStub()
	
	def forward(self, x):
		x = self.quant(x)
		...
		x = self.front3D(x)
		...
		x = self.front(x)
		x = self.dequant(x)		# 注意解量化的位置，后端已经被动态量化，所以无法走在QuantizeCPU中
		x, _ = self.back(x)
		x = self.fc(x)
		return x

注意解量化添加的位置是在后端网络之前，这是因为后端网络已经变成了DynamicQuantizedGRU，无法在QuantizeCPU中跑。
到这里，我们对这个网络的静态量化和动态量化的操作看起来是完成了。

3. 你可能会遇到的问题

（1）算子不支持

e.g. Runtime error: “MaxPool3d” , “add_()”, “mul()_” , “torch.zeros_like()” …QuantizeCPU…

这是由于静态量化不支持上述算子，解决办法之前提过，可以用解量化的思路。我被MaxPool3d恶心了好久，因为是sequential封装好的，所以我想了好久才用下面这种笨方法解决：

def forward(self, x):
	x = self.quant1(x)			# 模型最开始的量化
	for m in self.front3D:
	     if isinstance(m, nn.MaxPool3d):
	         x = self.dequant(x)
	         x = m(x)
	         x = self.quant2(x)		# 经过maxpool3d后再量化回来
	     else:
	         x = m(x)

这里需要注意的是，我使用了两个quant()，这是因为quant()中会保存scale, zero_point等信息，他是将FP32映射到int8的关键信息，所以模型内有多少次量化操作，就需要定义多少个QuantStub()，而DeQuantStub()只需定义一个就好，因为它内部没有储存什么信息。

对于原位操作add_()和mul()_，也可以用同样的思路解决。觉得麻烦也可以用torch.nn.quantized.FloatFunctional()，具体如下：

class basicblock(nn.Module):
	def __init__(self, ...):
		...
		self.func = torch.nn.quantized.FloatFunctional()
	def forward(self, x):
		...
		# out = out + residual
		out = self.func.add(out, residual)
		return out

e.g. Runtime error: "nai” , “torch.zeros_like()” …QuantizeCPU…

一些其他的算子，比如说batchnorm类，如果在模型出现单独的BN，而又无法融合，此时同样无法跑在QuantizedCPU上，需要解量化。

（2）模型精度下降很多

对于这个问题就比较玄学了，我认为的解决办法是对网络层逐个逐个的对比。那怎么做呢？
直接将不想量化的层的qconfig设置为None就行，对于本网络

model.front.layer1.qconfig = None
model.front.layer2.qconfig = None

经过这样的设置后，ResNet()中的layer1和layer2中的算子就没有被量化，所以在跑这些网络之前记得解量化。对于我的网络，消融后得到的结果，挺玄学的。

（3）对于这样的网络，可以单独量化吗？

我的探索是，对于单独动态量化，但是不能单独静态量化。如果对于CNN+RNN类的网络，单独静态量化的话，在observe阶段，会报错，大概是叫

Runtime error: “tuple” has no attribute “numel”…

这是因为RNN会有多个输出，模型无法直接观察，所以在本例中，我是先对模型做动态量化，再对模型做静态量化

最后简单总结一下我遇到的问题以及解决：

1. 对于不支持的算子、操作，直接解量化绕过，走FP32计算，也就是CPU后端
2. 如果有多个解量化操作，QuantStub()同样需要多个，而DeQuantStub()只需要定义一个。
3. 部分网络静态量化
4. tuple has no attribute numel
5. 精度下降明显