Cross-Layer Equalization跨层均衡

aimet解读

符合规则的模型结构

• 统一要求：单数据流，中间激活不流向其他地方
• 概念说明：
  • Conv: gruoups=1的普通卷积，包括TransposedConv和Conv
  • DepthwiseConv: 深度可分离卷积，groups=in_channels, in_channels=out_channels。
• cle均衡块：(相邻块的连接中间可能穿插Relu或者是Relu6等正缩放线性运算的算子)
  • Conv ==> Conv
  • Conv ==> DepthwiseConv ==> Conv
  • DepthwiseConv ==> Conv

优化目标

• 前提说明：

  • 正缩放线性运算函数特性：$$\begin{array}{r}f(sx)=sf(x)\end{array}$$，relu/relu6等算子符合该特性
  • 我们优化的函数总是以cle均衡块：Conv ==> Conv作为基础的情况，使用函数可表达为如下的函数
    $$\begin{array}{rl}h=& f(W^{(1)}x+b^{(1)})\\ y=& f(W^{(2)}h+b^{(2)})\\ y=& f(W^{(2)}f(W^{(1)}x+b^{(1)})+b^{(2)})\end{array}$$
  • 提出的优化目标和思路：
    • 出发点：cle的目的就是想要在模型推理结果不变的情况下，通过调整相连conv层的weight的每通道的权重，使得同一个weight的数值范围能够基本保持一致，这样能够让后续的per-layer量化效果能够和per-channel量化效果相当。
    • 思路：使用对角矩阵调整conv的每channel权重，使得该conv的每channel权重的数值范围range能够大致相同。而如何保证调整conv层权重后的模型推理运算结果不变，利用了正缩放线性运算函数特性。
  • 具体的权重调整公式如下所示：
    $$\begin{array}{rl}S=& diag(s_i)\\ h=& Sf(S^{-1}{W}^{(1)}x+S^{-1}{b}^{(1)})\\ y=& f(W^{(2)}Sf(S^{-1}{W}^{(1)}x+S^{-1}{b}^{(1)})+b^{(2)})\\ =& f({\tilde{W}}^{(2)}f({\tilde{W}}^{(1)}x+{\tilde{b}}^{(1)})+b^{(2)})\end{array}$$
  • 上述的推演公式可知如下的调整后的新的权重：
    • (9)公式：对应pre-layer（也就是第一个conv）的权重的调整，即对每output_channel上进行了对应的调整
    • (10)公式：对应cur-layer（也就是第二个conv）的权重的调整，即对每input_channel上进行了对应的调整
      $$\begin{array}{rl}{\tilde{W}}^{(1)}=& S^{-1}{W}^{(1)}\\ {\tilde{W}}^{(2)}=& W^{(2)}S\\ {\tilde{b}}^{(1)}=& S^{-1}{b}^{(1)}\end{array}$$
  • 具体的优化目标：

    • 理想情况下，对于同一个权重而言，希望每channel权重的range同整个权重的range相等。

    • 因此提出如下的优化目标：

      • ${\tilde{r}}_i^{(1)}$:表示每channel通道权重的数值范围（都是按照对称量化进行衡量的）
      • ${\tilde{R}}^{(1)}$:表示每整个权重的数值范围
      • (13)公式：最终的优化目标，获得一个$S$能够数值最大

      $$\begin{array}{r}{\tilde{p}}_i^{(1)}=\frac{{\tilde{r}}_i^{(1)}}{{\tilde{R}}^{(1)}}\\ \underset{S}{\max }\sum _i{\tilde{p}}_i^{(1)}{\tilde{p}}_i^{(2)}\end{array}$$

• 解优化目标：

  • 推导优化目标(基于(13)公式进行推导)
    $$\begin{array}{rl}{\tilde{r}}^{(1)}=& S^{-1}{r}^{(1)}\\ {\tilde{r}}^{(2)}=& r^{(2)}S\\ {\tilde{R}}^{(k)}=& {\max }_i({\tilde{r}}_i^{(k)})\\ \underset{S}{\max }\sum _i{\tilde{p}}_i^{(1)}{\tilde{p}}_i^{(2)}=& \underset{S}{\max }\sum _i\frac{{\tilde{r}}_i^{(1)}{\tilde{r}}_i^{(2)}}{{\tilde{R}}^{(1)}{\tilde{R}}^{(2)}}\\ =& \underset{S}{\max }\sum _i\frac{\frac{1}{s_i}{r}_i^{(1)}{s}_i{r}_i^{(2)}}{{\max }_j(\frac{1}{s_j}{r}_j^{(1)}){\max }_k(s_k{r}_k^{(2)})}\\ =& \sum _i{r}_i^{(1)}{r}_i^{(2)}\underset{S}{\max }\frac{1}{{\max }_j(\frac{1}{s_j}{r}_j^{(1)}){\max }_k(s_k{r}_k^{(2)})}\end{array}$$

  • 优化目标简化（基于(19)公式）

    • 这是一个最优化的问题，本来的优化目标从优化最大值变成了优化最小值的问题

    $$\begin{array}{rl}\underset{S}{\mathrm{argmax}}\sum _i{\tilde{p}}_i^{(1)}{\tilde{p}}_i^{(2)}=& \underset{S}{\mathrm{argmin}}[{\max }_j(\frac{1}{s_j}{r}_j^{(1)}){\max }_k(s_k{r}_k^{(2)})]\end{array}$$

  • 利用反证法，对于优化目标，可以推出如下的结论(26)：

    • 对于整个权重的range数值范围，我们有如下的假设，即pre-layer的权重的第$J$个channel有着最大的操作

    $$\begin{array}{rl}J=& \underset{j}{\mathrm{argmax}}(\frac{1}{s_j}{r}_j^{(1)})\\ K=& \underset{k}{\mathrm{argmax}}(s_k{r}_k^{(2)})\end{array}$$

    • 反证法证明过程：如果$J\ne K$，那么总是存在一个$\epsilon >0$，使得满足如下不等式。最为矛盾的是(25)行的不等式的含义，这个表明，还有更好的${\tilde{s}}_K$能够让(20)的优化目标更优的解。因此对于我们的优化目标来说，必然存在$J=K$

    $$\begin{array}{rl}{\tilde{s}}_K=& s_K-\epsilon \\ \frac{1}{s_J}{r}_J^{(1)}>& \frac{1}{{\tilde{s}}_K}{r}_K^{(1)}>\frac{1}{s_K}{r}_K^{(1)}\\ \frac{1}{s_J}{r}_J^{(1)}{s}_K{r}_K^{(2)}>& \frac{1}{s_J}{r}_J^{(1)}{\tilde{s}}_K{r}_K^{(2)}\end{array}$$

    $$\begin{array}{rl}\underset{j}{\mathrm{argmax}}(\frac{1}{s_j}{r}_j^{(1)})=& \underset{k}{\mathrm{argmax}}(s_k{r}_k^{(2)})\end{array}$$

  • 利用(26)的结论，如果我们再进一步简化优化目标，最后会发现优化目标变成一个跟$S$无关的解，而是依赖于$i=\underset{i}{\mathrm{argmax}}(r_i^{(1)}{r}_i^{(2)})$。因此这里为了能够解得所有的$S$，这里增加了如下所示的条件限制：
    $$\begin{array}{r}{\forall }_i:{\tilde{r}}_i^{(1)}={\tilde{r}}_i^{(2)}\end{array}$$

  • 因此结合(26)和(27)，我们可以得到满足条件的解如下：
    $$\begin{array}{rl}{s}_i=& \frac{1}{r_i^{(2)}}\sqrt{r_i^{(1)}{r}_i^{(2)}}\end{array}$$

限制和缺陷

从上述的优化过程中，我们可以看出，CLE存在一些缺陷和使用限制

• 优化推导过程中使用的是对称量化，因此cle后只适合对称量化权重的PTQ方式
• 因为pre-layer的bias改变的问题，改变了pre-layer的输出数值的范围，有可能出现每通道数值范围差别大，可能会导致后量化效果差。
• 对于pre-layer来说，是按照output_channel的方向进行权重微调，这个是符合per-channel的；但是对于cur-layer来说，是按照input_channel的方向进行权重微调，在含义上不符合per-channel的方式。

实操和结果显示

cle的实现函数

import numpy as np

def my_cle(pre_layer_weight, pre_layer_bias, cur_layer_weight):
    # 自建的cle函数，用于计算和
    channel_num = list(pre_layer_weight.shape)[0]
    r1 = np.array([np.abs(pre_layer_weight[i]).max() for i in range(channel_num)]).tolist()
    r2 = np.array([np.abs(cur_layer_weight[i]).max() for i in range(channel_num)]).tolist()
    scale = list()
    for a,b in zip(r1, r2):
        mul_v = a * b
        if mul_v == 0.0:
            scale.append(1.0)
        else:
            scale.append(math.sqrt(a/b))
    
    scale = np.array(scale, dtype=np.float32)
    pre_layer_weight_res = [pre_layer_weight[i]/s for i, s in enumerate(scale)]
    cur_layer_weight_res = [cur_layer_weight[i]*s for i, s in enumerate(scale)]
    pre_layer_bias_res   = [pre_layer_bias[i]/s for i, s in enumerate(scale)]
    pre_layer_weight_res = np.array(pre_layer_weight_res)
    cur_layer_weight_res = np.array(cur_layer_weight_res)
    pre_layer_bias_res = np.array(pre_layer_bias_res)
    return scale, pre_layer_weight_res, pre_layer_bias_res, cur_layer_weight_res

模型定义

import torch
class onnx2torch_cle(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5,
                                    stride=1, padding=2)
        self.relu1 = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(in_channels=16, out_channels=4, kernel_size=5,
                                    stride=1, padding=2)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu1(x)
        x = self.conv2(x)
        return x
    

整体的cle流程

def test_cle():
    model = onnx2torch_cle()
    weight_all(model)
    model = model.eval()
    model = model.cpu()
    model_aimet = copy.deepcopy(model)
    input_shape = [1, 3, 264, 264]
    dummy_input = torch.rand(input_shape)
    dummy_output = model(dummy_input)
    # 可视化权重
    output_png_root = os.path.join(os.path.dirname(__file__), 'doc')
    if not os.path.exists(output_png_root):
        os.mkdir(output_png_root)
    file_name = 'org_pre_layer_w'
    show_weight(model.conv.weight.detach().numpy(), file_name, os.path.join(output_png_root, '{}.png'.format(file_name)))
    file_name = 'org_cur_layer_w'
    show_weight(model.conv2.weight.detach().numpy().transpose(1,0,2,3), file_name, os.path.join(output_png_root, '{}.png'.format(file_name)))

    # cle微调权重
    scale, pre_layer_weight_res, pre_layer_bias_res, cur_layer_weight_res = my_cle(model.conv.weight.detach().numpy(), 
                                                                                   model.conv.bias.detach().numpy(), 
                                                                                   model.conv2.weight.detach().numpy().transpose(1,0,2,3))
    cur_layer_weight_res = torch.from_numpy(cur_layer_weight_res.transpose(1,0,2,3))
    pre_layer_weight_res = torch.from_numpy(pre_layer_weight_res)
    pre_layer_bias_res = torch.from_numpy(pre_layer_bias_res)
    model.conv.weight.data = pre_layer_weight_res
    model.conv.bias.data = pre_layer_bias_res
    model.conv2.weight.data = cur_layer_weight_res
    dummy_output_cle = model(dummy_input)
    a,b = compare_data(dummy_output.detach().numpy(), dummy_output_cle.detach().numpy())
    print(scale)
    print(a,b)
    file_name = 'cle_pre_layer_w'
    show_weight(model.conv.weight.detach().numpy(), file_name, os.path.join(output_png_root, '{}.png'.format(file_name)))
    file_name = 'cle_cur_layer_w'
    show_weight(model.conv2.weight.detach().numpy().transpose(1,0,2,3), file_name, os.path.join(output_png_root, '{}.png'.format(file_name)))
  
if __name__ == "__main__":
    test_cle()

测试结果

• 从compare的结果可以看出，cle前后的模型在同个输入下的输出是基本保持一致的(余弦相似度=1.0、相对最大误差=3.6033464e-07)。

• 下述的可视化结果如下，我们可以看到：

  • cle后model.conv.weight的每output_channel上的range相较cle前更加均衡
  • cle后model.conv2.weight的每input_channel上的range同model.conv.weight的每output_channel上的range基本相同

• 原始模型的权重分布如下：

  • model.conv.weight
    

  • model.conv2.weight
    

• cle后模型的权重分布如下：

  • model.conv.weight
    

  • model.conv2.weight
    

具体的测试脚本程序

可见下载链接