RepOptimizer: 其實(shí)是RepVGG2
前言
在神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)中，我們經(jīng)常會引入一些先驗(yàn)知識，比如ResNet的殘差結(jié)構(gòu)。然而我們還是用常規(guī)的優(yōu)化器去訓(xùn)練網(wǎng)絡(luò)。
在本工作中，我們提出將先驗(yàn)信息用于修改梯度數(shù)值，稱為梯度重參數(shù)化，對應(yīng)的優(yōu)化器稱為RepOptimizer。我們著重關(guān)注VGG式的直筒模型，訓(xùn)練得到RepOptVGG模型，他有著高訓(xùn)練效率，簡單直接的結(jié)構(gòu)和極快的推理速度。
官方倉庫：RepOptimizer
論文鏈接：Re-parameterizing Your Optimizers rather than Architectures
與RepVGG的區(qū)別
RepVGG加入了結(jié)構(gòu)先驗(yàn)（如1x1，identity分支），并使用常規(guī)優(yōu)化器訓(xùn)練。而RepOptVGG則是將這種先驗(yàn)知識加入到優(yōu)化器實(shí)現(xiàn)中 盡管RepVGG在推理階段可以把各分支融合，成為一個直筒模型。但是其訓(xùn)練過程中有著多條分支，需要更多顯存和訓(xùn)練時間。而RepOptVGG可是 真-直筒模型，從訓(xùn)練過程中就是一個VGG結(jié)構(gòu) 我們通過定制優(yōu)化器，實(shí)現(xiàn)了結(jié)構(gòu)重參數(shù)化和梯度重參數(shù)化的等價變換，這種變換是通用的，可以拓展到更多模型
將結(jié)構(gòu)先驗(yàn)知識引入優(yōu)化器
我們注意到一個現(xiàn)象，在特殊情況下，每個分支包含一個線性可訓(xùn)練參數(shù)，加一個常量縮放值，只要該縮放值設(shè)置合理，則模型性能依舊會很高。我們將這個網(wǎng)絡(luò)塊稱為Constant-Scale Linear Addition(CSLA)
我們先從一個簡單的CSLA示例入手，考慮一個輸入，經(jīng)過2個卷積分支+線性縮放，并加到一個輸出中：
我們考慮等價變換到一個分支內(nèi)，那等價變換對應(yīng)2個規(guī)則：
初始化規(guī)則
融合的權(quán)重需為：
更新規(guī)則
針對融合后的權(quán)重，其更新規(guī)則為：
這部分公式可以參考附錄A中，里面有詳細(xì)的推導(dǎo)
一個簡單的示例代碼為：
import torch
import numpy as np
np.random.seed(0)
np_x = np.random.randn(1, 1, 5, 5).astype(np.float32)
np_w1 = np.random.randn(1, 1, 3, 3).astype(np.float32)
np_w2 = np.random.randn(1, 1, 3, 3).astype(np.float32)
alpha1 = 1.0
alpha2 = 1.0
lr = 0.1
conv1 = torch.nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
conv2 = torch.nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
conv1.weight.data = torch.nn.Parameter(torch.tensor(np_w1))
conv2.weight.data = torch.nn.Parameter(torch.tensor(np_w2))
torch_x = torch.tensor(np_x, requires_grad=True)
out = alpha1 * conv1(torch_x) + alpha2 * conv2(torch_x)
loss = out.sum()
loss.backward()
torch_w1_updated = conv1.weight.detach().numpy() - conv1.weight.grad.numpy() * lr
torch_w2_updated = conv2.weight.detach().numpy() - conv2.weight.grad.numpy() * lr
print(torch_w1_updated + torch_w2_updated)
import torch
import numpy as np
np.random.seed(0)
np_x = np.random.randn(1, 1, 5, 5).astype(np.float32)
np_w1 = np.random.randn(1, 1, 3, 3).astype(np.float32)
np_w2 = np.random.randn(1, 1, 3, 3).astype(np.float32)
alpha1 = 1.0
alpha2 = 1.0
lr = 0.1
fused_conv = torch.nn.Conv2d(1, 1, kernel_size=3, padding=1, bias=False)
fused_conv.weight.data = torch.nn.Parameter(torch.tensor(alpha1 * np_w1 + alpha2 * np_w2))
torch_x = torch.tensor(np_x, requires_grad=True)
out = fused_conv(torch_x)
loss = out.sum()
loss.backward()
torch_fused_w_updated = fused_conv.weight.detach().numpy() - (alpha1**2 + alpha2**2) * fused_conv.weight.grad.numpy() * lr
print(torch_fused_w_updated)
在RepOptVGG中，對應(yīng)的CSLA塊則是將RepVGG塊中的3x3卷積，1x1卷積，bn層替換為帶可學(xué)習(xí)縮放參數(shù)的3x3卷積，1x1卷積
進(jìn)一步拓展到多分支中，假設(shè)s，t分別是3x3卷積，1x1卷積的縮放系數(shù)，那么對應(yīng)的更新規(guī)則為：
第一條公式對應(yīng)輸入通道==輸出通道，此時一共有3個分支，分別是identity，conv3x3， conv1x1
第二條公式對應(yīng)輸入通道！=輸出通道，此時只有conv3x3, conv1x1兩個分支
第三條公式對應(yīng)其他情況
需要注意的是CSLA沒有BN這種訓(xùn)練期間非線性算子(training-time nonlinearity)，也沒有非順序性(non sequential)可訓(xùn)練參數(shù)，CSLA在這里只是一個描述RepOptimizer的間接工具。
那么剩下一個問題，即如何確定這個縮放系數(shù)
HyperSearch
受DARTS啟發(fā)，我們將CSLA中的常數(shù)縮放系數(shù)，替換成可訓(xùn)練參數(shù)。在一個小數(shù)據(jù)集（如CIFAR100）上進(jìn)行訓(xùn)練，在小數(shù)據(jù)上訓(xùn)練完畢后，我們將這些可訓(xùn)練參數(shù)固定為常數(shù)。
具體的訓(xùn)練設(shè)置可參考論文
實(shí)驗(yàn)結(jié)果
實(shí)驗(yàn)效果看上去非常不錯，訓(xùn)練中沒有多分支，可訓(xùn)練的batchsize也能增大，模型吞吐量也提升不少。
在之前RepVGG中，不少人吐槽量化困難，那么在RepOptVGG下，這種直筒模型對于量化十分友好：
代碼簡單走讀
我們主要看 repoptvgg.py 這個文件，核心類是 RepVGGOptimizer
在reinitialize 方法中，它做的就是repvgg的工作，將1x1卷積權(quán)重和identity分支給融到3x3卷積中：
if len(scales) == 2:
conv3x3.weight.data = conv3x3.weight * scales[1].view(-1, 1, 1, 1) \
+ F.pad(kernel_1x1.weight, [1, 1, 1, 1]) * scales[0].view(-1, 1, 1, 1)
else:
assert len(scales) == 3
assert in_channels == out_channels
identity = torch.from_numpy(np.eye(out_channels, dtype=np.float32).reshape(out_channels, out_channels, 1, 1))
conv3x3.weight.data = conv3x3.weight * scales[2].view(-1, 1, 1, 1) + F.pad(kernel_1x1.weight, [1, 1, 1, 1]) * scales[1].view(-1, 1, 1, 1)
if use_identity_scales: # You may initialize the imaginary CSLA block with the trained identity_scale values. Makes almost no difference.
identity_scale_weight = scales[0]
conv3x3.weight.data += F.pad(identity * identity_scale_weight.view(-1, 1, 1, 1), [1, 1, 1, 1])
else:
conv3x3.weight.data += F.pad(identity, [1, 1, 1, 1])
然后我們再看下GradientMask生成邏輯，如果只有conv3x3和conv1x1兩個分支，根據(jù)前面的CSLA等價變換規(guī)則，conv3x3的mask對應(yīng)為：
mask = torch.ones_like(para) * (scales[1] ** 2).view(-1, 1, 1, 1)
而conv1x1的mask，需要乘上對應(yīng)縮放系數(shù)的平方，并加到conv3x3中間：
mask[:, :, 1:2, 1:2] += torch.ones(para.shape[0], para.shape[1], 1, 1) * (scales[0] ** 2).view(-1, 1, 1, 1)
如果還有Identity分支，我們則需要在對角線上加上1.0(Identity分支沒有可學(xué)習(xí)縮放系數(shù))
mask[ids, ids, 1:2, 1:2] += 1.0
如果有不明白Identity分支為什么對應(yīng)的是對角線，可以參考下筆者的圖解RepVGG
總結(jié)
這篇文章出來有段時間了，但是好像沒有很多人關(guān)注。在我看來這是個實(shí)用性很高的工作，解決了上一代RepVGG留下的小坑，真正實(shí)現(xiàn)了訓(xùn)練時完全直筒的模型，并且對量化，剪枝友好，十分適合實(shí)際部署。
