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 ) \begin{align} f(sx)=sf(x) \end{align} 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 ) ) \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} 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 ) ) \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} 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 ) \begin{align} \tilde{W}^{(1)}=&S^{-1}W^{(1)} \\ \tilde{W}^{(2)}=&W^{(2)}S \\ \tilde{b}^{(1)}=&S^{-1}b^{(1)} \end{align} W~(1)=W~(2)=b~(1)=S−1W(1)W(2)SS−1b(1)
- 具体的优化目标:
-
理想情况下,对于同一个权重而言,希望每channel权重的range同整个权重的range相等。
-
因此提出如下的优化目标:
- r ~ i ( 1 ) \tilde{r}^{(1)}_{i} r~i(1):表示每channel通道权重的数值范围(都是按照对称量化进行衡量的)
- R ~ ( 1 ) \tilde{R}^{(1)} R~(1):表示每整个权重的数值范围
- (13)公式:最终的优化目标,获得一个 S S S能够数值最大
p ~ i ( 1 ) = r ~ i ( 1 ) R ~ ( 1 ) max S ∑ i p ~ i ( 1 ) p ~ i ( 2 ) \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} p~i(1)=R~(1)r~i(1)Smaxi∑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 ) ) \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} r~(1)=r~(2)=R~(k)=Smaxi∑p~i(1)p~i(2)===S−1r(1)r(2)Smaxi(r~i(k))Smaxi∑R~(1)R~(2)r~i(1)r~i(2)Smaxi∑maxj(sj1rj(1))maxk(skrk(2))si1ri(1)siri(2)i∑ri(1)ri(2)Smaxmaxj(sj1rj(1))maxk(skrk(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 ) ) ] \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} Sargmaxi∑p~i(1)p~i(2)=Sargmin[maxj(sj1rj(1))maxk(skrk(2))]
-
利用反证法,对于优化目标,可以推出如下的结论(26):
- 对于整个权重的range数值范围,我们有如下的假设,即pre-layer的权重的第 J J J个channel有着最大的操作
J = arg max j ( 1 s j r j ( 1 ) ) K = arg max k ( s k r k ( 2 ) ) \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} J=K=jargmax(sj1rj(1))kargmax(skrk(2))
- 反证法证明过程:如果 J ≠ K J \neq K J=K,那么总是存在一个 ε > 0 \varepsilon > 0 ε>0,使得满足如下不等式。最为矛盾的是(25)行的不等式的含义,这个表明,还有更好的 s ~ K \tilde{s}_{K} s~K能够让(20)的优化目标更优的解。因此对于我们的优化目标来说,必然存在 J = K J = K 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 ) \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} s~K=sJ1rJ(1)>sJ1rJ(1)sKrK(2)>sK−εs~K1rK(1)>sK1rK(1)sJ1rJ(1)s~KrK(2)
arg max j ( 1 s j r j ( 1 ) ) = arg max k ( s k r k ( 2 ) ) \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} jargmax(sj1rj(1))=kargmax(skrk(2))
-
利用(26)的结论,如果我们再进一步简化优化目标,最后会发现优化目标变成一个跟 S S S无关的解,而是依赖于 i = arg max i ( r i ( 1 ) r i ( 2 ) ) i=\mathop{\arg\max}\limits_{i}(r^{(1)}_{i} r^{(2)}_{i}) i=iargmax(ri(1)ri(2))。因此这里为了能够解得所有的 S S S,这里增加了如下所示的条件限制:
∀ i : r ~ i ( 1 ) = r ~ i ( 2 ) \begin{align} {\forall}_i : \tilde{r}^{(1)}_{i}=\tilde{r}^{(2)}_{i} \end{align} ∀i:r~i(1)=r~i(2) -
因此结合(26)和(27),我们可以得到满足条件的解如下:
s i = 1 r i ( 2 ) r i ( 1 ) r i ( 2 ) \begin{align} s_{i} =& \frac{1}{r^{(2)}_{i}} \sqrt{r^{(1)}_{i} r^{(2)}_{i}} \end{align} si=ri(2)1ri(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 npdef 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 * bif 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_resmodel.conv.bias.data = pre_layer_bias_resmodel.conv2.weight.data = cur_layer_weight_resdummy_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
-
具体的测试脚本程序
可见下载链接