神经网络量化----TensorRT深刻解读
目录
前言
本文将聚焦于英伟达TensorRT训练后量化的算法。
论文地址为:https://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf
代码地址为:官方好像没有公布代码,可以参考的有https://github.com/deepglint/EasyQuant,https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py
一、TensorRT简介
整个量化算法使用对称量化(-max, max)-> (-127, 127)
权重的量化所需的最值直接统计得出,激活值的量化使用饱和映射的方式,设置最佳的阈值作为最值来进行量化。
计算最佳阈值的方法:1. 统计激活值的直方图,2. 采用遍历的方法找到量化后KL散度最小时对应的最佳阈值。伪代码如下:
二、难点
1.架构
以上所给的参考代码语言分别为caffe和mxnet,caffe代码里有保存和加载文件的操作,对于量化后的参数不能直观地和网络架构联系在一起,mxnet代码里的功能较全,但是语言比较小众化,开发新的算法难度较大。
故计划采用pytorch来进行量化参数的计算和保存,且跟网络是一个整体,新建module即可,比较方便。
2.功能
有的卷积后有激活和BN,所以还要将其考虑在内。
所以在处理之前需要将网络中的卷积、BN、ReLU融合在一起。
在量化时,不考虑是否有ReLU,全部量化在(-max, max)之间。(格林深瞳算法只考虑了带有ReLU的,即将(0,max)量化到(0,127)),这样就简化了运算,不用再分情况了。
三、实现
1.conv和ReLU的融合
from torch import nn
import torch
# the module that replace BN layer
class DummyModule(nn.Module):
def __init__(self):
super(DummyModule, self).__init__()
def forward(self, x):
return x
# BN flod
def bn_folding(conv, bn):
# ******************** BN parameter *********************
mean = bn.running_mean
std = torch.sqrt(bn.running_var + bn.eps)
gamma = bn.weight
beta = bn.bias
# ******************* conv parameter********************
w = conv.weight
w_fold = w.clone()
if conv.bias is not None:
b = conv.bias
else:
b = mean.new_zeros(mean.shape)
b_fold = b.clone()
w_fold = w * (gamma / std).reshape([conv.out_channels, 1, 1, 1])
b_fold = beta + (b - mean) * (gamma / std)
bnfold_conv = nn.Conv2d(conv.in_channels,
conv.out_channels,
conv.kernel_size,
conv.stride,
conv.padding,
groups=conv.groups,
bias=True)
bnfold_conv.weight.data = w_fold
bnfold_conv.bias.data = b_fold
return bnfold_conv
'''BN must be after convolution'''
def model_bn_folding(model):
children = list(model.named_children())
# children = list(model.named_modules())
#print(children)
name_temp = None
child_temp = None
for name, child in children:
#print(name, ' ', child)
if isinstance(child, nn.BatchNorm2d):
bnfold_conv = bn_folding(child_temp, child) # BN融合
model._modules[name_temp] = bnfold_conv
model._modules[name] = DummyModule()
child_temp = None
elif isinstance(child, nn.Conv2d):
name_temp = name
child_temp = child
else:
model_bn_folding(child)
return model
2.conv和ReLU的融合
新建一个module将卷积和ReLU包含在内了。
import torch
from torch import nn
import torch.nn.functional as F
from quant_utils import ConvRelu, LinearRelu, DummyModule
# device = torch.device("cpu")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
''