上一篇介紹了RepVGG，RepVGG存在量化問題，Repopt通過將先驗融入優化器中，統一訓練與測試模型解決了其量化不友好的問題。

論文鏈接： Re-parameterizing Your Optimizers rather than Architectures

Introduction

Repopt提出將模型結構的先驗信息直接用于修改梯度數值，其稱為梯度重參數化，對應的優化器稱為RepOptimizer。Repopt著重關注VGG式的直筒模型，訓練得到RepOptVGG模型與VGG結構一致，有著高訓練效率，簡單直接的結構和極快的推理速度。

與RepVGG的不同：
1）RepVGG在訓練過程中加入了結構先驗（shortcut，1x1 branch），在推理時，將多支路融合成單路3x3卷積。而RepOptVGG將結構先驗轉移至梯度中，通過設計的RepOpt優化器實現。
2）在結構上，RepOptVGG是真-直筒結構，模型在訓練與測試時保持一致。RepVGG訓練時存在多支路需要更多的顯存與訓練時間。
3）RepOptVGG通過定制優化器，實現了結構重參與梯度重參的等效變化。

Idea

Repopt發現結構先驗的一個有趣現象：當每個分支只包含一個線性可訓練算子，如果正確設置常尺度值，模型的性能會提高。我們將這種線性塊稱為Constant Scale Linear Addition（CSLA）。我們可以用單個算子替換一個CSLA塊，并通過設計優化器改變梯度實現等價的訓練動態。Repopt將這種乘數稱為Grad Mult，如上圖所示。

證明：用常規的SGD訓練一個CSLA塊相當于用修改的梯度訓練一個簡單的卷積

CSLA塊中每個分支只包含一個可訓練線性算子，并且結構中不存在BN或者dropout等非線性操作。Repopt發現用常規的SGD訓練一個CSLA塊相當于用修改的梯度訓練一個簡單的卷積。下面用一個簡單的例子證明這個結論。

假設CSLA由兩個相同形狀的卷積組成，其中每個核包含一個可訓練線性算子。如下面公式所示，其中${\alpha }_A,{\alpha }_B$為可訓練線性算子，W為卷積的參數，X是輸入，Y為CSLA的輸出，*表示卷積操作。

對應的梯度重參公式$Y_{GR}=X?W^{\prime }$,其中$W^{\prime }$表示梯度重參后的卷積，假設損失函數為L，訓練迭代數為i，卷積參數W的梯度表示為$\frac{?L}{?W}$,$F(\frac{?L}{?W^{\prime }})$表示對應梯度重參上的任意變化，我們希望通過數次訓練后CSLA的輸出與梯度重參后的輸出一致，即

通過卷積的線性可加性，我們需要保證公式6

在i=0迭代開始前，正確的初始化確保了公式6的等價性，初始化如公式7所示

下面，我們用數學歸納法證明在$W^{\prime }$的梯度上進行適當的變換后，公式6的等價性始終成立，W梯度更新的公式如下

更新相應的CSLA塊，我們獲得公式10

我們使用$F(\frac{?L}{?W^{\prime }})$來更新$W^{\prime }$，這就意味著

假設在迭代第i次時，公式6，10，11成立，那么可以獲得公式12

對公式6取偏導數，有公式13

我們獲得等式14，即$F(\frac{?L}{?W^{\prime }})$的準確形式

由公式11，14，我們可以推到出，當迭代到i+1次時，下面等式成立

由于假設公式6成立

通過初始條件公式7，8，以及數學歸納法我們可以證明當i>=0時，公式6成立。同時，我們知道$F(\frac{?L}{?W^{\prime }})$的準確形式，如公式14所示。

Method

上文，已經介紹了Repopt找到一個合適的結構先驗CSLA塊，并通過數學歸納證明可以通過梯度重參將CSLA等效為簡單的卷積操作，下面，我們使用RepOpt-VGG作為展示例，具體介紹Repopt如何設計和描述梯度重參的行為。

在RepOptVGG中，對應的CSLA塊則是將RepVGG塊中的3x3卷積，1x1卷積，bn層替換為帶可學習縮放參數的3x3卷積，1x1卷積。進一步拓展到多分支中，假設s，t分別是3x3卷積，1x1卷積的縮放系數，那么對應的更新規則為：

對公式3的理解需要結合RepVGG，當輸入與輸出通道不等時，只存在conv3x3, conv1x1兩個分支，其中conv1x1可以等效為特殊的conv3x3，因此梯度可以重參為$s_c^2+t_c^2$，如上文所證明一樣。而當輸入與輸出通道相等時，此時一共有3個分支，分別是identity，conv3x3， conv1x1，Identity也可以等效為特殊的conv3x3，其卷積核由0，1組成，所以梯度重參為$1+s_c^2+t_c^2$。

需要注意的是CSLA沒有BN這種訓練期間非線性算子(training-time nonlinearity)，也沒有非順序性(non sequential)可訓練參數，CSLA在這里只是一個描述RepOptimizer的間接工具。

那么剩下一個問題，即如何確定這個縮放系數

HyperSearch

受DARTS啟發，我們將CSLA中的常數縮放系數，替換成可訓練參數。在一個小數據集（如CIFAR100）上進行訓練，在小數據上訓練完畢后，我們將這些可訓練參數固定為常數。

Code

LinearAddBlock定義的是CSLA塊，該模塊只在確定HyperSearch的時候被訓練。

class LinearAddBlock(nn.Module):def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1,dilation=1, groups=1, padding_mode='zeros', use_se=False, is_csla=False, conv_scale_init=1.0):super(LinearAddBlock, self).__init__()self.in_channels = in_channelsself.relu = nn.ReLU()self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)self.scale_conv = ScaleLayer(num_features=out_channels, use_bias=False, scale_init=conv_scale_init)self.conv_1x1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, bias=False)self.scale_1x1 = ScaleLayer(num_features=out_channels, use_bias=False, scale_init=conv_scale_init)if in_channels == out_channels and stride == 1:self.scale_identity = ScaleLayer(num_features=out_channels, use_bias=False, scale_init=1.0)self.bn = nn.BatchNorm2d(out_channels)if is_csla: # Make them constantself.scale_1x1.requires_grad_(False)self.scale_conv.requires_grad_(False)if use_se:raise NotImplementedError("se block not supported yet")else:self.se = nn.Identity()def forward(self, inputs):out = self.scale_conv(self.conv(inputs)) + self.scale_1x1(self.conv_1x1(inputs))if hasattr(self, 'scale_identity'):out += self.scale_identity(inputs)out = self.relu(self.se(self.bn(out)))return outclass ScaleLayer(torch.nn.Module):def __init__(self, num_features, use_bias=True, scale_init=1.0):super(ScaleLayer, self).__init__()self.weight = Parameter(torch.Tensor(num_features))init.constant_(self.weight, scale_init)self.num_features = num_featuresif use_bias:self.bias = Parameter(torch.Tensor(num_features))init.zeros_(self.bias)else:self.bias = Nonedef forward(self, inputs):if self.bias is None:return inputs * self.weight.view(1, self.num_features, 1, 1)else:return inputs * self.weight.view(1, self.num_features, 1, 1) + self.bias.view(1, self.num_features, 1, 1)

RealVGGBlock是RepOptVGG的真實模塊，結構簡單如下所示。

class RealVGGBlock(nn.Module):def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1,dilation=1, groups=1, padding_mode='zeros', use_se=False,):super(RealVGGBlock, self).__init__()self.relu = nn.ReLU()self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False)self.bn = nn.BatchNorm2d(out_channels)if use_se:raise NotImplementedError("se block not supported yet")else:self.se = nn.Identity()def forward(self, inputs):out = self.relu(self.se(self.bn(self.conv(inputs))))return out

假設我們已經通過小數據訓練獲得了HyperSearch需要的scales，那么在訓練RepOptVGG時，RepVGGOptimizer需要在初始化時候將CSLA塊的scales賦值給RealVGGBlock，賦值的過程如reinitialize所示，對應了Method中的公式3。

def reinitialize(self, scales_by_idx, conv3x3_by_idx, use_identity_scales):for scales, conv3x3 in zip(scales_by_idx, conv3x3_by_idx):in_channels = conv3x3.in_channelsout_channels = conv3x3.out_channelskernel_1x1 = nn.Conv2d(in_channels, out_channels, 1, device=conv3x3.weight.device)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) == 3assert in_channels == out_channelsidentity = torch.from_numpy(np.eye(out_channels, dtype=np.float32).reshape(out_channels, out_channels, 1, 1)).to(conv3x3.weight.device)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])

我們在梯度重參過程中需要獲取梯度Mask，與reinitialize過程相似分為3種情況，具體實現如下所示。

def generate_gradient_masks(self, scales_by_idx, conv3x3_by_idx, cpu_mode=False):self.grad_mask_map = {}for scales, conv3x3 in zip(scales_by_idx, conv3x3_by_idx):para = conv3x3.weightif len(scales) == 2:mask = torch.ones_like(para, device=scales[0].device) * (scales[1] ** 2).view(-1, 1, 1, 1)mask[:, :, 1:2, 1:2] += torch.ones(para.shape[0], para.shape[1], 1, 1, device=scales[0].device) * (scales[0] ** 2).view(-1, 1, 1, 1)else:mask = torch.ones_like(para, device=scales[0].device) * (scales[2] ** 2).view(-1, 1, 1, 1)mask[:, :, 1:2, 1:2] += torch.ones(para.shape[0], para.shape[1], 1, 1, device=scales[0].device) * (scales[1] ** 2).view(-1, 1, 1, 1)ids = np.arange(para.shape[1])assert para.shape[1] == para.shape[0]mask[ids, ids, 1:2, 1:2] += 1.0if cpu_mode:self.grad_mask_map[para] = maskelse:self.grad_mask_map[para] = mask.cuda()

通過Repopt梯度重參的方式將結構先驗轉化為梯度先驗，可以統一訓練與測試模型結構，有效解決RepVGG量化不友好問題，其結構在YOLOV6中被使用，并表現出極佳的性能。

總結

以上是生活随笔為你收集整理的Repoptimizer论文理解与代码分析的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。