【PTQ】Cross-Layer Equalization跨层均衡-证明和实践详细解读

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{align} f(sx)=sf(x) \end{align}$ ，relu/relu6等算子符合该特性
- 我们优化的函数总是以cle均衡块：Conv ==> Conv作为基础的情况，使用函数可表达为如下的函数
  $\begin{align} 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{align}$
- 提出的优化目标和思路：
  - 出发点：cle的目的就是想要在模型推理结果不变的情况下，通过调整相连conv层的weight的每通道的权重，使得同一个weight的数值范围能够基本保持一致，这样能够让后续的per-layer量化效果能够和per-channel量化效果相当。
  - 思路：使用对角矩阵调整conv的每channel权重，使得该conv的每channel权重的数值范围range能够大致相同。而如何保证调整conv层权重后的模型推理运算结果不变，利用了正缩放线性运算函数特性。
- 具体的权重调整公式如下所示：
  $\begin{align} 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{align}$
- 上述的推演公式可知如下的调整后的新的权重：
  - (9)公式：对应pre-layer（也就是第一个conv）的权重的调整，即对每output_channel上进行了对应的调整
  - (10)公式：对应cur-layer（也就是第二个conv）的权重的调整，即对每input_channel上进行了对应的调整
    $\begin{align} \tilde{W}^{(1)}=&S^{-1}W^{(1)} \\ \tilde{W}^{(2)}=&W^{(2)}S \\ \tilde{b}^{(1)}=&S^{-1}b^{(1)} \end{align}$
- 具体的优化目标：
  - 理想情况下，对于同一个权重而言，希望每channel权重的range同整个权重的range相等。
  - 因此提出如下的优化目标：
    - $\tilde{r}^{(1)}_{i}$ :表示每channel通道权重的数值范围（都是按照对称量化进行衡量的）
    - $\tilde{R}^{(1)}$ :表示每整个权重的数值范围
    - (13)公式：最终的优化目标，获得一个 $S$ 能够数值最大
    $\begin{align} \tilde{p}^{(1)}_{i}=\frac{\tilde{r}^{(1)}_{i}}{\tilde{R}^{(1)}} \\ \mathop{\max}\limits_{S} \sum\limits_{i} \tilde{p}^{(1)}_{i} \tilde{p}^{(2)}_{i} \end{align}$
解优化目标：
- 推导优化目标(基于(13)公式进行推导)
  $\begin{align} \tilde{r}^{(1)}=&S^{-1}r^{(1)} \\ \tilde{r}^{(2)}=&r^{(2)}S \\ \tilde{R}^{(k)}=&\mathop{\max}_{i}(\tilde{r}^{(k)}_{i}) \\ \mathop{\max}\limits_{S} \sum\limits_{i} \tilde{p}^{(1)}_{i} \tilde{p}^{(2)}_{i} = &\mathop{\max}\limits_{S} \sum\limits_{i} \frac{\tilde{r}^{(1)}_{i}\tilde{r}^{(2)}_{i}}{\tilde{R}^{(1)}\tilde{R}^{(2)}} \\ =&\mathop{\max}\limits_{S} \sum\limits_{i} \frac{\frac{1}{s_{i}}r^{(1)}_{i} s_{i} r^{(2)}_{i}}{\mathop{\max}_{j}(\frac{1}{s_{j}}r^{(1)}_{j}) \mathop{\max}_{k}(s_{k} r^{(2)}_{k})} \\ =& \sum\limits_{i} r^{(1)}_{i} r^{(2)}_{i} \mathop{\max}\limits_{S} \frac{1}{\mathop{\max}_{j}(\frac{1}{s_{j}}r^{(1)}_{j}) \mathop{\max}_{k}(s_{k} r^{(2)}_{k})} \end{align}$
- 优化目标简化（基于(19)公式）
  - 这是一个最优化的问题，本来的优化目标从优化最大值变成了优化最小值的问题
  $\begin{align} \mathop{\arg\max}\limits_{S} \sum\limits_{i} \tilde{p}^{(1)}_{i} \tilde{p}^{(2)}_{i}=&\mathop{\arg\min}\limits_{S} [\mathop{\max}_{j}(\frac{1}{s_{j}}r^{(1)}_{j}) \mathop{\max}_{k}(s_{k} r^{(2)}_{k})] \end{align}$
- 利用反证法，对于优化目标，可以推出如下的结论(26)：
  - 对于整个权重的range数值范围，我们有如下的假设，即pre-layer的权重的第 $J$ 个channel有着最大的操作
  $\begin{align} J=&\mathop{\arg\max}\limits_{j}(\frac{1}{s_{j}}r^{(1)}_{j}) \\ K=&\mathop{\arg\max}\limits_{k}(s_{k} r^{(2)}_{k}) \end{align}$
  - 反证法证明过程：如果 $\neq K$ ，那么总是存在一个 $\varepsilon > 0$ ，使得满足如下不等式。最为矛盾的是(25)行的不等式的含义，这个表明，还有更好的 $\tilde{s}_{K}$ 能够让(20)的优化目标更优的解。因此对于我们的优化目标来说，必然存在 $J = K$
  $\begin{align} \tilde{s}_{K} =& s_{K} - \varepsilon \\ \frac{1}{s_{J}}r^{(1)}_{J} >& \frac{1}{\tilde{s}_{K}}r^{(1)}_{K} > \frac{1}{s_{K}}r^{(1)}_{K} \\ \frac{1}{s_{J}}r^{(1)}_{J} s_{K} r^{(2)}_{K} >& \frac{1}{s_{J}}r^{(1)}_{J} \tilde{s}_{K} r^{(2)}_{K} \end{align}$
  
  $\begin{align} \mathop{\arg\max}\limits_{j}(\frac{1}{s_{j}}r^{(1)}_{j})=&\mathop{\arg\max}\limits_{k}(s_{k} r^{(2)}_{k}) \end{align}$
- 利用(26)的结论，如果我们再进一步简化优化目标，最后会发现优化目标变成一个跟 $S$ 无关的解，而是依赖于 $i=\mathop{\arg\max}\limits_{i}(r^{(1)}_{i} r^{(2)}_{i})$ 。因此这里为了能够解得所有的 $S$ ，这里增加了如下所示的条件限制：
  $\begin{align} {\forall}_i : \tilde{r}^{(1)}_{i}=\tilde{r}^{(2)}_{i} \end{align}$
- 因此结合(26)和(27)，我们可以得到满足条件的解如下：
  $\begin{align} s_{i} =& \frac{1}{r^{(2)}_{i}} \sqrt{r^{(1)}_{i} r^{(2)}_{i}} \end{align}$

限制和缺陷

从上述的优化过程中，我们可以看出，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