04-模型量化

库的介绍

现在写的这么丑陋，具体的等什么时候要上库了再说

Observer

Observer负责根据指定的方式监测Tensor的max_val和min_val，譬如激活值常用滑动平均Observer，权重值则往往直接根据当前数值直接获取。Observer可按layer或channel的粒度进行监测。

Quantizer

Quantizer负责对Tensor进行量化，Quantizer包含了Observer并根据Observer的监测结果确定scale等量化参数。现有的Quantizer只分了对称量化器和无符号量化器（适用relu后的激活值量化），非对称量化和非线性量化（K-Means那种）都没写。

量化的权重层

目前写了4个层：QConv2d、QLinear、QConv2d_BN、QLinear_BN，它们分别继承了nn.Conv2d、nn.Linear（没错1d和3d我都没写），用于实现量化的特征提取。

代码实现

逐通道Observer

逐层Observer很好写就不多说了，整个Tensor拉max和min后赋值或者做指数平滑平均就行，逐层Observer的max和min是一个值（Tensor.Size([])），scale也是一个值，量化时直接“input / scale”就行。
逐通道Observer会遇到几个麻烦一些的问题：

1. torch.max()、torch.min()只支持取整个Tensor的最值，或者指定一个维度上取最值，最笨的写法需要写成：

minval = torch.min(torch.min(torch.min(input, 3, keepdim=True)[0], 2, keepdim=True)[0], 1, keepdim=True)[0]

1. 得到的maxval、minval以及scale的shape是（Tensor.Size([64])）这样的，而不像逐层的时候是一个值，量化时如果直接“input / scale”无法进行广播，shape为[128, 64, 32, 32]的Tensor不能和shape为[64]的Tensor进行运算，只能和shape为[1, 64, 1, 1]的Tensor进行运算；而对于全连接层，shape为[128, 512]的Tensor同样需要和shape为[1, 512]的Tensor进行运算。

代码：

    def _per_channel_min_max(self, x):
        shape = [1] * x.ndimension()
        shape[self.ch_axis] = x.shape[self.ch_axis]            # 确定min、max（scale也是）的shape
        new_axis_list = list(range(x.ndimension()))
        new_axis_list[self.ch_axis] = 0
        new_axis_list[0] = self.ch_axis
        y = x.permute(tuple(new_axis_list))                    # 把通道这一维度换到最前面来
        y = y.flatten(start_dim=1)                            # 拉平除了通道的所有维度
        min_vals = torch.min(y, 1)[0].reshape(shape)
        max_vals = torch.max(y, 1)[0].reshape(shape)

函数一开始就确定好scale需要的shape，省的到最后一次次unsqueeze。
接下来的Tensor拉平后求最值的方法仿照了pytorch官方的Observer写法，（以卷积为例）将128643232的Tensor利用permute函数转置为641283232后，利用flatten函数拉平从dim=1开始的所有维度得到128*131072的Tensor后，对dim=1求最大最小值后reshape为需要的shape。

Quantizer

在模型量化中我们已经讨论过需要为反传构建“直通估计器”，否则梯度会被截平，因此我们需要自己改写量化的梯度反传过程，使其直接返回输入的梯度值。pytorch里直接自己写一个backward函数似乎没有用，百度以后得出了以下有效的写法：

from torch.autograd import Function
class Round(Function):
    @staticmethod
    def forward(self, x):
        output = torch.round(x)
        return output
    @staticmethod
    def backward(self, grad_output):
        grad_input = grad_output.clone()
        return grad_input
class Quantizer(nn.Module):
    def round(self, input):
        out = Round.apply(input)
        return out

量化的权重层

以QConv2d_BN为例简单介绍一下量化层：

class QConv2d_BN(nn.Conv2d):
    ...
    def forward(self, x):
        if self.training and self.use_running_status_to_train is False:
            with torch.no_grad():
                tmp = self._conv(x, self.weight)
            batch_mean = torch.mean(tmp, dim=[0, 2, 3])
            batch_var = torch.var(tmp, dim=[0, 2, 3])
            self._update_running_status(batch_mean, batch_var)
            weight, bias = self._calc_para_from_statistic(batch_mean, batch_var)
        else:
            weight, bias = self._calc_para_from_statistic(self.running_mean, self.running_var)
        analog_qx = self.act_quantizer(x)
        analog_qw = self.weight_quantizer(weight)
        output = self._conv(analog_qx, analog_qw, bias)
        if self.act is not None:
            output = self.act(output)
        return output

该层实现了对输入激活值与权重的模拟量化，并进行常规特征提取操作。“use_running_status_to_train”flag可用于控制是否使用running_mean和running_var进行训练。
如果使用统计量进行训练，第一次卷积获取batch_mean和batch_var时，必须加with torch.no_grad()，不然main里面loss.backward()会报错：
RuntimeError: Trying to backward through the graph a second time,# but the saved intermediate results have already been freed. Specify retain_graph=True when calling# backward the first time.

获取需要量化的层

使用了正则表达式，来自FAIR的开源代码，真的是黑科技。。。用正则表达式用户可以非常方便地给定哪些层需要进行量化。

def get_layers(model, filter_regexp):
    """
    Filters out the layers according to a regexp. Note that
    we omit biases.
    Args:
        - model: a nn.Module
        - filter_regexp: a regexp to filter the layers to keep
          according to their name in model.named_parameters().
          For instance, the regexp:
             down_layers\\.[123456]\\.(conv[12]|identity\\.conv))
          is keeping blocks down_layers from 1 to 6, and inside
          each block is keeping conv1, conv2 and identity.conv.
    Remarks:
        - We add (module\\.)? at the beginning of the regexp to
          account for the possible use of nn.parallel.DataParallel
    """
    # named_patameters()是一个返回tuple(name: str, parameter: Tensor)的generator
    # itemgetter(0)用于获取每个tuple的第一个元素，即名字
    all_layers = map(itemgetter(0), model.named_parameters())
    # 我们需要的是层的名字而非参数的名字，去除包含bias的参数，保证每个层中只留下了一个key
    all_layers = list(filter(lambda x: "bias" not in x, all_layers))
    # 把 “.weight” 去掉以获取层的名字(or .weight_orig for spectral norm)
    all_layers = map(lambda x: x.replace(".weight_orig", ""), all_layers)
    all_layers = map(lambda x: x.replace(".weight", ""), all_layers)
    # return filtered layers
    filter_regexp = "(module\\.)?" + "(" + filter_regexp + ")"
    r = re.compile(filter_regexp)
    mod_list = list(filter(r.match, all_layers))
    return mod_list

简单讲讲正则表达式：
由于“.”在正则表达式中代表任何单个字符，需要表示实际的“.”需要写“.”，可是python里反斜杠是字符串的转义符，要表示实际的“\”需要写“\”，于是在正则表达式里为了表示“.”需要打“\.”。
[123456]表示取里面的任意一个字符，[1-3]?[0-9]表示0~39
(downsample|conv[12])表示downsample和conv1和conv2
代码中“(module\.)?”表示在正则表达式开头加上module或不加module都匹配，用于DataParallel模型，正则表达式中“?”表示可以出现1次或0次。
“features.”，表示以features开头的任何字符串都匹配，“”表示前一项可以匹配任意次数。

取出这些层后经过一系列打包操作，将conv或fc转换为合成的量化层，剩下的全部做idendity层，这涉及到如何根据名字取出对应层并将它们替换。由于一般层的名字都像“features.0.conv1”这样，我们不能在model里直接找“features.0.conv1”这个属性，而需要以“.”作为分割逐层寻找子属性。
来自pytorch官方的写法：

def _get_module(model, submodule_key):
    tokens = submodule_key.split('.')
    cur_mod = model
    for s in tokens:
        cur_mod = getattr(cur_mod, s)
    return cur_mod
def _set_module(model, module, submodule_key):
    tokens = submodule_key.split('.')
    sub_tokens = tokens[:-1]
    cur_mod = model
    for s in sub_tokens:
        cur_mod = getattr(cur_mod, s)
    setattr(cur_mod, tokens[-1], module)

量化部分代码冗余

本来以为python支持多重继承，我可以自己写一个class _QuantLayer，然后在QConv2d_BN里同时继承nn.Conv2d和_QuantLayer，但是后来发现无法对多个基类进行初始化，python的多重继承是假的。。。
如果只继承nn.Conv2d和nn.Linear这些pytorch的权重层，那就每个class都要写一遍冗余的量化部分的代码，很烦。郑老师的意见是把量化部分的代码抽成函数，直接做函数调用。虽然感觉也不是特别优雅的写法不过姑且先这样吧。

踩的一系列坑

层转换时的参数复制

千万别漏复制了啊。。。可是又不能用state_dict因为命名不一样，一开始就是全连接层的只复制了bias没复制weight我真的是血妈爆炸……
用load_state_dict方法和一个个Tensor手动赋值没有区别，不需要担心load_state_dict做了额外的事情。Tensor的单纯赋值只需要：a[:] = b就可以，不需要考虑深浅拷贝问题。如果是Parameter的赋值，直接这么赋值可能会报错操作了non-leaf Tensor，我也不知道什么意思，只需要改成：a.data[:] = b就可以。

eps的取值

这里涉及的eps主要又两个——第一个是BN层的eps，第二个Quantizer的eps。
BN层的eps主要用于防止FeatureMap在除了sqrt(var)后由于var是0而直接爆炸，换了我们融合BN的卷积层就是用于防止计算融合BN的weight时，weight大的直接爆炸。这个eps根据nn.BatchNorm2d一样取1e-5没有问题，取太小的话weight还是比较容易过大的。
Quantizer的eps就不能随便取了，其主要作用是用于防止Qx = (x - zeropoint) / scale时，因scale过小而爆炸，但是如果量化的位数比较多（8位甚至8位以上），有的时候scale可能本来就只有1e-5这么大，这时候加上一个1e-5这么大的eps就会影响量化过程，所以设置为1e-7甚至1e-9更为合适。

nn.BatchNorm的momentum

pytorch的BN的momentum默认为0.1，但是这个值是反过来的！它是new 0.1 + old 0.9，复制到新的层的时候需要用1去减原来的momentum。

实验结果

根据不同的量化位数需要使用不同的学习率和策略。
这里使用的是ResNet26_64网络，分4个stage每个stage有3个BasicBlock，数据集为CIFAR100。我们在训练得到的float网络的基础上进行重训练量化。batchsize=100，SGD(wd=5e-5, momentum=0.9)

1. 监测激活值的指数平滑平均的Observer的momentum需要设置为0.99，设为0.9多掉半个点。
2. 低比特必须全程用running_mean running_var跑，不能用batch_mean batch_var，否则训练过程会非常不稳定甚至爆炸，最终结果多掉3个多点。
3. 最后一层全连接可以量化（只多掉0.1个点），但是第一层绝对不能量化（我没试但一定会炸）

神奇的发现

CIFAR100模型，宽度为64的ResNet，第一个最接近输入的卷积层，64个通道有40个是全0！冗余度非常非常高！问了下郑老师这种情况只在浅层出现，深层的冗余不会冗余到这种程度。自己做实验也是，第一个stage的ResNet Block里就已经很少有全0通道了，第二个stage就完全没了。