【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

优化目标

• 前提说明：

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

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

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

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

      p ~ i ( 1 ) = r ~ i ( 1 ) R ~ ( 1 ) max ⁡ S ∑ i p ~ i ( 1 ) p ~ i ( 2 )

      p~i(1)=r~i(1)R~(1)maxS⁡∑ip~i(1)p~i(2)" role="presentation" style="position: relative;">p~(1)i=r~(1)iR~(1)maxS∑ip~(1)ip~(2)i$$\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}$$
      p~​i(1)​=R~(1)r~i(1)​​Smax​i∑​p~​i(1)​p~​i(2)​​​

• 解优化目标：

  • 推导优化目标(基于(13)公式进行推导)
    r ~ ( 1 ) = S − 1 r ( 1 ) r ~ ( 2 ) = r ( 2 ) S R ~ ( k ) = max ⁡ i ( r ~ i ( k ) ) max ⁡ S ∑ i p ~ i ( 1 ) p ~ i ( 2 ) = max ⁡ S ∑ i r ~ i ( 1 ) r ~ i ( 2 ) R ~ ( 1 ) R ~ ( 2 ) = max ⁡ S ∑ i 1 s i r i ( 1 ) s i r i ( 2 ) max ⁡ j ( 1 s j r j ( 1 ) ) max ⁡ k ( s k r k ( 2 ) ) = ∑ i r i ( 1 ) r i ( 2 ) max ⁡ S 1 max ⁡ j ( 1 s j r j ( 1 ) ) max ⁡ k ( s k r k ( 2 ) )

    r~(1)=S−1r(1)r~(2)=r(2)SR~(k)=maxi⁡(r~i(k))maxS⁡∑ip~i(1)p~i(2)=maxS⁡∑ir~i(1)r~i(2)R~(1)R~(2)=maxS⁡∑i1siri(1)siri(2)maxj⁡(1sjrj(1))maxk⁡(skrk(2))=∑iri(1)ri(2)maxS⁡1maxj⁡(1sjrj(1))maxk⁡(skrk(2))" role="presentation" style="position: relative;">r~(1)=r~(2)=R~(k)=maxS∑ip~(1)ip~(2)i===S−1r(1)r(2)Smaxi(r~(k)i)maxS∑ir~(1)ir~(2)iR~(1)R~(2)maxS∑i1sir(1)isir(2)imaxj(1sjr(1)j)maxk(skr(2)k)∑ir(1)ir(2)imaxS1maxj(1sjr(1)j)maxk(skr(2)k)$$\begin{array}{rl}{\tilde{r}}^{(1)}=& S^{-1}{r}^{(1)}\\ {\tilde{r}}^{(2)}=& r^{(2)}S\\ {\tilde{R}}^{(k)}=& \underset{i}{max}({\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}$$
    r~(1)=r~(2)=R~(k)=Smax​i∑​p~​i(1)​p~​i(2)​===​S−1r(1)r(2)Smaxi​(r~i(k)​)Smax​i∑​R~(1)R~(2)r~i(1)​r~i(2)​​Smax​i∑​maxj​(sj​1​rj(1)​)maxk​(sk​rk(2)​)si​1​ri(1)​si​ri(2)​​i∑​ri(1)​ri(2)​Smax​maxj​(sj​1​rj(1)​)maxk​(sk​rk(2)​)1​​​

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

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

    arg ⁡ max ⁡ S ∑ i p ~ i ( 1 ) p ~ i ( 2 ) = arg ⁡ min ⁡ S [ max ⁡ j ( 1 s j r j ( 1 ) ) max ⁡ k ( s k r k ( 2 ) ) ]

    arg⁡maxS⁡∑ip~i(1)p~i(2)=arg⁡minS⁡[maxj⁡(1sjrj(1))maxk⁡(skrk(2))]" role="presentation" style="position: relative;">argmaxS∑ip~(1)ip~(2)i=argminS[maxj(1sjr(1)j)maxk(skr(2)k)]$$\begin{array}{rl}\underset{S}{\arg max}\sum _i{\tilde{p}}_i^{(1)}{\tilde{p}}_i^{(2)}=& \underset{S}{\arg min}[\underset{j}{max}(\frac{1}{s_j}{r}_j^{(1)})\underset{k}{max}(s_k{r}_k^{(2)})]\end{array}$$
    Sargmax​i∑​p~​i(1)​p~​i(2)​=​Sargmin​[maxj​(sj​1​rj(1)​)maxk​(sk​rk(2)​)]​​

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

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

    J = arg ⁡ max ⁡ j ( 1 s j r j ( 1 ) ) K = arg ⁡ max ⁡ k ( s k r k ( 2 ) )

    J=arg⁡maxj⁡(1sjrj(1))K=arg⁡maxk⁡(skrk(2))" role="presentation" style="position: relative;">J=K=argmaxj(1sjr(1)j)argmaxk(skr(2)k)$$\begin{array}{rl}J=& \underset{j}{\arg max}(\frac{1}{s_j}{r}_j^{(1)})\\ K=& \underset{k}{\arg max}(s_k{r}_k^{(2)})\end{array}$$
    J=K=​jargmax​(sj​1​rj(1)​)kargmax​(sk​rk(2)​)​​
    • 反证法证明过程：如果$J\ne K$，那么总是存在一个$\epsilon >0$，使得满足如下不等式。最为矛盾的是(25)行的不等式的含义，这个表明，还有更好的${\tilde{s}}_K$能够让(20)的优化目标更优的解。因此对于我们的优化目标来说，必然存在$J=K$

    s ~ K = s K − ε 1 s J r J ( 1 ) > 1 s ~ K r K ( 1 ) > 1 s K r K ( 1 ) 1 s J r J ( 1 ) s K r K ( 2 ) > 1 s J r J ( 1 ) s ~ K r K ( 2 )

    s~K=sK−ε1sJrJ(1)>1s~KrK(1)>1sKrK(1)1sJrJ(1)sKrK(2)>1sJrJ(1)s~KrK(2)" role="presentation" style="position: relative;">s~K=1sJr(1)J>1sJr(1)JsKr(2)K>sK−ε1s~Kr(1)K>1sKr(1)K1sJr(1)Js~Kr(2)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}$$
    s~K​=sJ​1​rJ(1)​>sJ​1​rJ(1)​sK​rK(2)​>​sK​−εs~K​1​rK(1)​>sK​1​rK(1)​sJ​1​rJ(1)​s~K​rK(2)​​​

    arg ⁡ max ⁡ j ( 1 s j r j ( 1 ) ) = arg ⁡ max ⁡ k ( s k r k ( 2 ) )

    arg⁡maxj⁡(1sjrj(1))=arg⁡maxk⁡(skrk(2))" role="presentation" style="position: relative;">argmaxj(1sjr(1)j)=argmaxk(skr(2)k)$$\begin{array}{rl}\underset{j}{\arg max}(\frac{1}{s_j}{r}_j^{(1)})=& \underset{k}{\arg max}(s_k{r}_k^{(2)})\end{array}$$
    jargmax​(sj​1​rj(1)​)=​kargmax​(sk​rk(2)​)​​

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

    ∀i:r~i(1)=r~i(2)" role="presentation" style="position: relative;">∀i:r~(1)i=r~(2)i$$\begin{array}{r}{\mathrm{\forall }}_i:{\tilde{r}}_i^{(1)}={\tilde{r}}_i^{(2)}\end{array}$$
    ∀i​:r~i(1)​=r~i(2)​​​

  • 因此结合(26)和(27)，我们可以得到满足条件的解如下：
    s i = 1 r i ( 2 ) r i ( 1 ) r i ( 2 )

    si=1ri(2)ri(1)ri(2)" role="presentation" style="position: relative;">si=1r(2)ir(1)ir(2)i−−−−−−√$$\begin{array}{rl}{s}_i=& \frac{1}{r_i^{(2)}}\sqrt{r_i^{(1)}{r}_i^{(2)}}\end{array}$$
    si​=​ri(2)​1​ri(1)​ri(2)​ ​​​

限制和缺陷

从上述的优化过程中，我们可以看出，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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

模型定义

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
    
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

整体的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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

测试结果

• 从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
    

具体的测试脚本程序

可见下载链接