首发：AIWalker
作者:HappyAIWalker最近Happy在尝试进行图像超分的INT8量化，发现：pytorch量化里面的坑真多，远不如TensorFlow的量化好用。不过花了点时间终于还是用pytorch把图像超分模型完成了量化，以EDSR为例，模型大小73%,推理速度提升40%左右(PC端)，视觉效果几乎无损，定量指标待补充。有感于网络上介绍量化的博客一堆，但真正有帮助的较少，所以Happy会尽量以图像超分为例提供一个完整的可复现的量化示例。
在前面的文章中，笔者对Pytorch的“Post Training Static Quantization，PTSQ”进行了原理性的介绍。接下来，我们将以EDSR这个图像超分网络为例进行说明。
准备工作
在真正开始量化之前，我们需要准备好要进行量化的模型，本文以EDSR-baseline模型为基础进行。所以大家可以直接下载官方预训练模型，EDSR的Pytorch官方实现code连接如下：
github.com/thstkdgus35/EDSR-PyTorch EDSRx4-baseline预训练模型下载连接如下：
https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x4-6b446fab.pt 除了要准备上述预训练模型与code外，我们还需要准备校验数据，在这里笔者采用的DIV2K数据，该数据集下载链接如下：
https://cv.snu.ac.kr/research/EDSR/DIV2K.tar 模型转换
正如上一篇文章所介绍的，在量化之前需要对模型进行op融合操作，而EDSR官方的实现code是对于融合操作是不太方便的，所以笔者对EDSR进行了一些实现上的调整。调整成如下形式(注：这里的实现code部分参数写成了固定参数)：
class ResBlock(nn.Module):
def __init__(self, channels=64):
super(ResBlock, self).__init__()
self.conv1 = nn.Conv2d(channels, channels, 3, 1, 1)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(channels, channels, 3, 1, 1)
def forward(self, x):
identity = x
conv1 = self.conv1(x)
relu = self.relu(conv1)
conv2 = self.conv2(relu)
output = conv2 + identity
return output
class EDSR(nn.Module):
def __init__(self,
num_blocks=16,
num_features=64,
block=ResBlock):
super(EDSR, self).__init__()
self.head = nn.Conv2d(3, num_features, 3, 1, 1)
body = [
block(num_features) for _ in range(num_blocks)
]
body.append(nn.Conv2d(num_features, num_features, 3, 1, 1))
self.body = nn.Sequential(*body)
self.tail = nn.Sequential(
nn.Conv2d(num_features, num_features * 4, 3, 1, 1),
nn.PixelShuffle(upscale_factor=2),
nn.Conv2d(num_features, num_features * 4, 3, 1, 1),
nn.PixelShuffle(upscale_factor=2),
nn.Conv2d(num_features, 3, 3, 1, 1)
)
def forward(self, x, **kwargs):
x = self.head(x)
res = self.body(x)
res += x
x = self.tail(res)
return x 也许有同学会说，模型转换后原始的预训练模型还能导入吗？直接导入肯定是不行的，checkpoint的key发生了变化，所以我们需要对下载的checkpoint进行一下简单的转换。checkpoint的转换code如下（注：这些转换可以都是写死的，已经确认过的）：
checkpoint = torch.load("edsr_baseline_x4-6b446fab.pt", map_location='cpu')
newStateDict = OrderedDict()
for key, val in checkpoint.items():
if 'head' in key:
newStateDict[key.replace('.0.', '.')] = val
elif 'mean' in key:
continue
# newStateDict[key] = val
elif 'tail' in key:
if '.0.0.' in key:
newStateDict[key.replace('.0.0.', '.0.')] = val
elif '.0.2.' in key:
newStateDict[key.replace('.0.2.', '.2.')] = val
else:
newStateDict[key.replace('.1.', '.4.')] = val
elif 'body' in key:
if '.body.0.' in key:
newStateDict[key.replace(".body.0.", '.conv1.')] = val
elif '.body.2.' in key:
newStateDict[key.replace(".body.2.", '.conv2.')] = val
elif "16" in key:
newStateDict[key] = val
torch.save(newStateDict, "edsr-baseline-fp32.pth.tar") 对比原始code的同学应该会发现：EDSR中的add\_mean与sub\_mean不见了。是的，笔者将add\_mean与sub\_mean移到了网络外面，不对其进行量化，具体为什么这样做，见后面的介绍。
除了上述操作外，我们还需要提供前述EDSR实现的量化版本模型，这个没太多需要介绍的，直接看code(主要体现在三点：插入量化节点(即QuantStub与DequantStub)、add转换(即FloatFunctional)、fuse\_model模块(即fuse\_model函数))：
class QuantizableResBlock(ResBlock):
def __init__(self, *args, **kwargs):
super(QuantizableResBlock, self).__init__(*args, **kwargs)
self.add = FloatFunctional()
def forward(self, x):
identity = x
conv1 = self.conv1(x)
relu = self.relu(conv1)
conv2 = self.conv2(relu)
output = self.add.add(identity, conv2)
return output
def fuse_model(self):
fuse_modules(self, ['conv1', 'relu'], inplace=True)
class QuantizableEDSR(EDSR):
def __init__(self, *args, **kwargs):
super(QuantizableEDSR, self).__init__(*args, **kwargs)
self.quant = QuantStub()
self.dequant = DeQuantStub()
self.add = FloatFunctional()
def forward(self, x):
x = self.quant(x)
x = self.head(x)
res = self.body(x)
res = self.add.add(res, x)
x = self.tail(res)
x = self.dequant(x)
return x
def fuse_model(self):
for m in self.modules():
if type(m) == QuantizableResBlock:
m.fuse_model() 模型量化
在上一篇文章中，我们也介绍了PTSQ的几个步骤（额外包含了模型的构建与保存）。
- init: 模型的定义、预训练模型加载、inplace操作替换为非inplace操作；
- config：定义量化时的配置方式，这里以
fbgemm为例，它的activation量化方式为Historam，weight量化方式为per\_channel； - fuse：模型中的op融合，比如相邻的Conv+ReLU融合，Conv+BN+ReLU融合等等；
- prepare: 量化前的准备工作，也就是对每个需要进行量化的op插入Observer；
- feed: 送入校验数据，前面插入的Observer会针对这些数据进行量化前的信息统计;
- convert：用于在将非量化op转换成量化op，比如将nn.Conv2d转换成nnq.Conv2d， 同时会根据Observer所观测的信息进行nnq.Conv2d中的量化参数的统计，包含scale、zero\_point、qweight等；
- save：用于保存量化好的模型参数.
Init
模型的创建与预训练模型，这个比较简单了，直接上code（注：PTSQ模式下模型应当是eval模式）。
checkpoint = torch.load("edsrx4-baseline-fp32.pth.tar")
model = QuantizableEDSR(block=QuantizableResBlock)
model.load_state_dict(checkpoint)
_replace_relu(model)
model.eval() config
这个步骤主要是为了指定与推理引擎搭配的一些量化方式，比如X86平台应该采用fbgemm方式进行量化，而ARM平台则应当采用qnnpack方式量化。本文主要是在PC端进行，所以选择了fbgemm进行，相关配置信息如下：
backend = 'fbgemm'
torch.backends.quantized.engine = backend
model.qconfig = torch.quantization.QConfig( activation=default_histogram_observer, weight=default_per_channel_weight_observer
) Fuse&Prepare
Fuse与Prepare两个步骤的作用主要是
- 进行OP的融合，比如Conv+ReLU的融合，Conv+BN+ReLU的融合，这个可以见前述实现code中的'fuse\_model'，pytorch目前提供了几种类型的融合。我们只需知道就可以了，这块不用太过关心，两行code就可以完成：
model.fuse_model()
torch.quantization.prepare(model, inplace=True) - 插入Observer，在每个需要进行量化的op中插入Observer，不同的量化方式会有不同的Observer，它将对喂入的校验数据进行统计，比如统计数据的最大值、最小值、直方图分布等等。
Feed
这个步骤需要采用校验数据喂入到上述准备好的模型中，这个就比较简单了，按照常规模型的测试方式处理就可以了，参考code如下：
注：笔者这里用了100张数据，这个用全部也可以，不过耗时会更长
meanBGR = torch.FloatTensor((0.4488, 0.4371, 0.4040)).view(3, 1, 1) * 255
data_root = "${DIV2K_train_LR_bicubic/X4}"
for index in range(1, 100):
image_path = os.path.join(data_root, f"{index:04d}.png")
inputs = preprocess(image_path)
inputs -= meanBGR
with torch.no_grad():
output = model(inputs) Convert&Save
在完成前面几个步骤后，我们就可以将浮点类型的模型进行量化了，这个只需要一行code就可以。在转换过程中，它会将nn.Conv2d这类浮点类型op转换成量化版op：nnq.Conv2d。
torch.quantization.convert(model, inplace=True)
torch.save(model.state_dict(), "edsrx4-baseline-qint8.pth.tar") 经过上面的几个步骤，我们就完成了EDSR模型的INT8量化，也将其进行了保存。也就是说完成了初步的量化工作，因为接下来的测试论证很关键，如果量化损失很严重也不行的。
量化模型测试
接下来，我们对上述量化好的模型进行一下测试看看效果。量化模型的调用code如下(与常规模型的调用有一点点的区别)：
def fp32edsr(block=ResBlock, pretrained=None):
model = EDSR(block=block)
if pretrained:
state_dict = torch.load(pretrained, map_location="cpu")
model.load_state_dict(state_dict)
return model
def qint8edsr(block=QuantizableResBlock, pretrained=None, quantize=False):
model = QuantizableEDSR(block=block)
_replace_relu(model)
if quantize:
backend = 'fbgemm'
quantize_model(model, backend)
else:
assert pretrained in [True, False]
if pretrained:
state_dict = torch.load(pretrained, map_location="cpu")
model.load_state_dict(state_dict)
return model
def quantize_model(model, backend):
if backend not in torch.backends.quantized.supported_engines:
raise RuntimeError("Quantized backend not supported ")
torch.backends.quantized.engine = backend
model.eval()
_dummy_input_data = torch.rand(1, 3, 64, 64)
# Make sure that weight qconfig matches that of the serialized models
if backend == 'fbgemm':
model.qconfig = torch.quantization.QConfig(
activation=torch.quantization.default_histogram_observer,
weight=torch.quantization.default_per_channel_weight_observer)
elif backend == 'qnnpack':
model.qconfig = torch.quantization.QConfig(
activation=torch.quantization.default_histogram_observer,
weight=torch.quantization.default_weight_observer)
model.fuse_model()
torch.quantization.prepare(model, inplace=True)
model(_dummy_input_data)
torch.quantization.convert(model, inplace=True) 从上面code可以看到：相比fp32模型，量化模型多了两步骤：
- replace=True的op替换为replace=False的op；
- 模型的最简单量化版本，完成初步的op替换。
结合上述code，我们就可以直接对DIV2K数据进行测试了，测试的部分code摘录如下：
index = 1
image_path = os.path.join(data_root, f"{index:04d}.png")
inputs = preprocess(image_path)
inputs -= meanBGR
with torch.no_grad():
output1 = model(inputs)
output2 = fmodel(inputs)
output1 += meanBGR
output2 += meanBGR
show1 = post_process(output1)
cv2.imwrite(f"results/{index:03d}-init8.png", show1)
show2 = post_process(output2)
cv2.imwrite(f"results/{index:03d}-fp32.png", show2) 
上图给出了DIV2K训练集中0016的两种模型的效果对比，左图为FP32模型的超分效果，右图为INT8量化模型的超分效果。可以看到：量化后模型在效果上是视觉无损的(就是说：量化损失导致的效果下降不可感知)。总而言之，量化前后模型的对比可以参考下表（PC端测试，测试数据为DIV2K，速度为平均速度）。
注意事项
- 为什么要将add\_mean与sub\_mean移到网络外面不参与量化呢？
从我们的量化对比来看，将其移到外面效果更佳。可能也跟add\_mean与sub\_mean中的参数有关，两者只是简单的均值处理， 这个地方的量化会导致weight值出现较大偏差，进而影响后续的量化精度。
- 在量化方式方面，该如何选择呢？
在量化方式方面，activation支持：HistogramObserver，MinMaxObserver，， weight支持：PerChannelMinMaxObserver，MinMaxObserver. 从我们的量化对比来看，Histogram+PerChannelMinMax这种组合要比MinMaxObserver+PerChannelMinMax更佳。下图给出了DIV2K训练集中0018数据采用第二种量化组合效果对比，可以感知到明显的量化损失。
推荐阅读
本文章著作权归作者所有，任何形式的转载都请注明出处。更多动态滤波，图像质量，超分辨相关请关注我的专栏深度学习从入门到精通。
