Intel Distiller工具包-量化实现3

 本系列文章

Intel Distiller工具包-量化实现1

Intel Distiller工具包-量化实现2

Intel Distiller工具包-量化实现3


回顾

  •  上面文章中介绍了Distiller及Quantizer基类,后训练量化器;基类定义了重要的变量,如replacement_factory(dict,用于记录待量化module对应的wrapper);此外定义了量化流程,包括 预处理(BN折叠,激活优化等)、量化模块替换、后处理 等主要步骤;后训练量化器则在基类的基础上实现了后训练量化的功能;
  •  本文继续介绍继承自Quantizer的子类量化器,包括
    • PostTrainLinearQuantizer(前文)
    • QuantAwareTrainRangeLinearQuantizer(本文)
    • PACTQuantizer(后续)
    • NCFQuantAwareTrainQuantizer(后续)
  • 本文代码也挺多的,由于没法全部贴出来,有些地方说的不清楚的,还请读者去参考源码;

QuantAwareTrainRangeLinearQuantizer

  • 量化感知训练量化器;将量化过程插入模型代码中,对模型进行训练;该过程使得模型参数对对量化过程有所拟合,所以最终得到的模型的效果 一般要比 后训练量化模型要好一些;
  • QuantAwareTrainRangeLinearQuantizer的类定义如下:可以看到比后训练量化器的定义简单不少;Intel Distiller工具包-量化实现3_第1张图片
  • 构造函数:前面都是检查和默认设置;核心是红框中 对 参数、激活值 设置的 量化感知 方式;Intel Distiller工具包-量化实现3_第2张图片
  • activation_replace_fn:这是对激活值做量化感知的实现方式,和前文后训练量化使用的 模块替换方式 一样,即返回一个 wrapper,这里是 FakeQuantizationWrapperIntel Distiller工具包-量化实现3_第3张图片
  • FakeQuantizationWrapper:定义如下,forward中输入先经过原module计算(得到的是原来的激活输出),然后对输出(下一个module的输入)做伪量化(fake_q);Intel Distiller工具包-量化实现3_第4张图片
  • FakeLinearQuantization:定义如下,该module做的事是 对输入做伪量化;具体细节包括 训练过程确定 激活值的范围并更新scale、zp(infer则直接使用训练过程最后的scale、zp);使用 LinearQuantizeSTE(straight-through-estimator) 实现伪量化;
    class FakeLinearQuantization(nn.Module):
        def __init__(self, num_bits=8, mode=LinearQuantMode.SYMMETRIC, ema_decay=0.999, dequantize=True, inplace=False):
            """
    
            :param num_bits:
            :param mode:
            :param ema_decay: 激活值范围使用EMA进行跟踪
            :param dequantize:
            :param inplace:
            """
            super(FakeLinearQuantization, self).__init__()
    
            self.num_bits = num_bits
            self.mode = mode
            self.dequantize = dequantize
            self.inplace = inplace
    
            # We track activations ranges with exponential moving average, as proposed by Jacob et al., 2017
            # https://arxiv.org/abs/1712.05877(激活值范围使用EMA进行跟踪)
            # We perform bias correction on the EMA, so we keep both unbiased and biased values and the iterations count
            # For a simple discussion of this see here:
            # https://www.coursera.org/lecture/deep-neural-network/bias-correction-in-exponentially-weighted-averages-XjuhD
            self.register_buffer('ema_decay', torch.tensor(ema_decay))  # 设置buffer,buffer用于非参的存储,会存于model state_dict
            self.register_buffer('tracked_min_biased', torch.zeros(1))
            self.register_buffer('tracked_min', torch.zeros(1))  # 保存无偏值
            self.register_buffer('tracked_max_biased', torch.zeros(1))  # 保存有偏值
            self.register_buffer('tracked_max', torch.zeros(1))
            self.register_buffer('iter_count', torch.zeros(1))  # 保存迭代次数
            self.register_buffer('scale', torch.ones(1))
            self.register_buffer('zero_point', torch.zeros(1))
    
        def forward(self, input):
            # We update the tracked stats only in training
            #
            # Due to the way DataParallel works, we perform all updates in-place so the "main" device retains
            # its updates. (see https://pytorch.org/docs/stable/nn.html#dataparallel)
            # However, as it is now, the in-place update of iter_count causes an error when doing
            # back-prop with multiple GPUs, claiming a variable required for gradient calculation has been modified
            # in-place. Not clear why, since it's not used in any calculations that keep a gradient.
            # It works fine with a single GPU. TODO: Debug...
            if self.training:  # 训练阶段要收集收据
                with torch.no_grad():
                    current_min, current_max = get_tensor_min_max(input)  # input是激活函数输出值
                self.iter_count += 1
                # 有偏值为正常加权值,无偏值为 有偏值/(1-decay**step)
                self.tracked_min_biased.data, self.tracked_min.data = update_ema(self.tracked_min_biased.data,
                                                                                 current_min, self.ema_decay,
                                                                                 self.iter_count)
                self.tracked_max_biased.data, self.tracked_max.data = update_ema(self.tracked_max_biased.data,
                                                                                 current_max, self.ema_decay,
                                                                                 self.iter_count)
    
            if self.mode == LinearQuantMode.SYMMETRIC:
                max_abs = max(abs(self.tracked_min), abs(self.tracked_max))
                actual_min, actual_max = -max_abs, max_abs
                if self.training:  # 激活值的范围数值经EMA更新后需要重新计算scale和zp
                    self.scale.data, self.zero_point.data = symmetric_linear_quantization_params(self.num_bits, max_abs)
            else:
                actual_min, actual_max = self.tracked_min, self.tracked_max
                signed = self.mode == LinearQuantMode.ASYMMETRIC_SIGNED
                if self.training:  # 激活值的范围数值经EMA更新后需要重新计算scale和zp
                    self.scale.data, self.zero_point.data = asymmetric_linear_quantization_params(self.num_bits,
                                                                                                  self.tracked_min,
                                                                                                  self.tracked_max,
                                                                                                  signed=signed)
    
            input = clamp(input, actual_min.item(), actual_max.item(), False)
            # 执行量化、反量化操作,并且该过程无需额外梯度
            input = LinearQuantizeSTE.apply(input, self.scale, self.zero_point, self.dequantize, False)
    
            return input
  • LinearQuantizeSTE:这是实现伪量化的核心,定义如下;它被定义为 torch.autograd.Function,指定了如何反向传播(STE方式)Intel Distiller工具包-量化实现3_第5张图片
  • 接下来看一下对参数的量化感知实现(linear_quantize_param),直接使用LinearQuantizeSTEIntel Distiller工具包-量化实现3_第6张图片
  • 注:distiller的量化感知训练量化器中虽然定义了如何 对参数做量化感知训练,但是并没有使用,有点奇怪;

总结

  • 本文介绍了distiller量化器基类Quantizer的一个子类:PostTrainLinearQuantizer;
  • 核心部分是 激活值、参数值的 量化感知训练的实现;对激活值量化感知的实现还是采用wrapper的方式,对参数则是直接使用STE;具体细节包括了 FakeQuantizationWrapper、FakeLinearQuantization、LinearQuantizeSTE;

你可能感兴趣的:(技术问题,#,深度学习,#,量化,机器学习,人工智能)