神经网络量化----TensorRT深刻解读

神经网络量化----TensorRT深刻解读

---

 

目录

神经网络量化----TensorRT深刻解读

前言

一、TensorRT简介

二、难点

1.架构

2.功能

三、实现

1.conv和ReLU的融合

2.conv和ReLU的融合

 quant_utils.py

3.调用示例

总结

---

---

前言

本文将聚焦于英伟达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")

''