3 本文方法
3.1 操作集合
1 Simple=>Fast
最近许多多分支架构的理论FLOPs比VGG要低，但运行起来可能并不会更快。例如，VGG-16的FLOPs是EfficientNet-B3的8.4倍，但在1080Ti上运行速度要快1.8倍，这意味着前者的计算密度是后者的15倍。除了Winograd conv带来的加速外，内存访问成本(MAC)和并行度这两个重要因素对速度有很大的影响，但内存访问成本和并行度并没有被计算在内。
例如，虽然所需要的分支加法或连接的计算是微不足道的，但MAC是重要的。此外，MAC在分组卷积中占据了很大一部分时间。另一方面，在相同的 FLOPs情况下，具有高并行度的模型可能比另一个具有低并行度的模型要快得多。由于多分支拓扑在初始化和自动生成的体系结构中被广泛采用，因此使用了多个小的运算符而不是几个大的运算符。
之前有工作说明了分割操作的数量(即在一个构建块中单个Conv或池化操作的数量)NASNET-A是13，这对GPU等拥有强大并行计算能力的设备不太友好，并引入了额外的开销，如内核启动和同步。相比之下，这个数字在ResNets中是2或3，本文所提的方法将其设为1:仅仅是单个Conv。
2 Memory-economical
多分支拓扑的内存效率是非常低下的，因为每个分支的结果都需要保留到叠加或连接，这大大提高了内存占用。如下图所示，Residual Block的输入需要保持到加法为止。假设Block保持feature map大小，则内存占用保持在2倍作为输入空间。相比之下，普通拓扑允许特定层的输入所占用的内存在操作完成后立即释放。
3 Flexible
多分支拓扑对体系结构规范施加了约束。例如，ResNet要求将conv层组织为Residual Block，这限制了灵活性，因为每个Residual Block的最后一个conv层必须产生相同形状的张量，否则Shortcut将没办法进行。
更糟糕的是，多分支拓扑限制了通道剪枝的应用。通道剪枝是一种去除一些不重要通道的实用技术，有些方法可以通过自动发现每一层合适的宽度来优化模型结构。然而，多分支模型使修剪变得棘手，并导致显著的性能下降或较低的加速比。相比之下，普通架构允许我们根据需求自由配置每个conv层，并进行修剪，以获得更好的性能-效率平衡。
3.2 多分支结构的训练时间
简单的卷积网络有很多优点，但有一个致命的缺点:性能差。例如，VGG-16使用BN等组件后在ImageNet上可以达到72%的top-1精度，但是现在看来是比较低了。本文所提结构受ResNet re-parameterization方法，构造一个shortcut分支模型，并使用一个Residual Block学习。当和的尺寸不匹配，它变成了, 是一个卷积实现的shortcut。
由于多分支拓扑在推理方面存在缺陷，但分支似乎对训练有益，因此本文使用多个分支对多个模型进行单独的训练时间集成。为了使大多数members更浅或更简单，作者使用类似了resnet的identity(仅当维度匹配时)和分支以便构建Block的训练信息流为。
本文所提架构只是简单地堆叠几个这样的块来构建训练模型。从相同的角度来看，模型是由个members和n个这样的块组成的集合。训练后，将其等价转换为，其中h由一个单独的卷积层实现，其参数由训练后的参数通过一系列代数推导而来。
3.3 重新指定推理时间模型的参数
这里描述了如何将一个训练块转换成一个单独的3×3的conv层进行推理。注意，在加法之前的每个分支中都使用了BN，如下图：
形式上，用来表示具有输入通道和输出通道的 conv层的核，用表示分支的核。使用分别表示conv后均值,标准差和学习因子以及BN层的偏差，分别表示conv后均值,标准差和学习因子以及BN层的偏差，和为 identity branch。设、分别为输入和输出，为卷积算子。如果，，有：
否则，如果不使用identity branch，因此上述方程只有前两项。这里bn是推理时间bn函数：
首先将每一个BN及其前面的卷积层转换成一个带有bias向量的卷积。让是从转换过来的kernel和bias:
上述变换也适用于identity branch，因为单位映射可以被视为以单位矩阵为核的1x1 conv。如Figure 4所示，经过这样的变换，将得到一个3x3 kernel，2个1x1 kernel，和3个bias向量。然后获得通过添加了3个bias向量的最终bias，最后3x3 kernel通过添加1x1 kernel的中心点上3x3的kernel，这可以很容易地由第1个零填充两个1x1内核实现3x3和添加3个kernel。需要注意的是，这些转换的等价性要求3x3层和1x1层具有相同的stride，而1x1层的padding配置要比3x3层少一个像素。例如，对于一个3x3层，填充一个像素的输入，这是最常见的情况，1x1层应该有padding=0。
3.4 架构说明
下表显示了RepVGG的规格，包括深度和宽度。RepVGG是VGG风格的，它采用简单的拓扑结构，大量使用3×3 conv，但它不像VGG那样使用最大池化，因为希望主体只有一种操作类型。这里安排3×3 layer分为5个阶段，第一个阶段的第1层下降样品的stride=2。对于图像分类，使用全局平均池化，然后使用全连接层作为head。对于其他任务，特定于任务的head可以用于任何层产生的特性。
根据3个简单的原则来决定每个阶段的层数：
- 第1阶段的操作分辨率较大，耗时较长，因此只使用一Layer以降低延迟。
- 最后1个阶段应该有更多的通道，所以只使用一个Layer来保存参数。
- 将大部分Layer放入最后的第2阶段(ImageNet输出分辨率为14x14)，然后是ResNet及其最新版本(例如，ResNet-101在其14x14分辨率阶段使用了69层)。
这里让这5个阶段分别有1、2、4、14、1个层来构建一个名为RepVGG-A的实例。还构建了一个更深层的RepVGG-B，在阶段2、3和4中有更多的层，等等。
import torch.nn as nn import numpy as np import torch def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1): result = nn.Sequential() result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False)) result.add_module('bn', nn.BatchNorm2d(num_features=out_channels)) return result class RepVGGBlock(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', deploy=False): super(RepVGGBlock, self).__init__() self.deploy = deploy self.groups = groups self.in_channels = in_channels assert kernel_size == 3 assert padding == 1 padding_11 = padding - kernel_size // 2 self.nonlinearity = nn.ReLU() if deploy: self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode) else: self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups) self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups) print('RepVGG Block, identity = ', self.rbr_identity) def forward(self, inputs): if hasattr(self, 'rbr_reparam'): return self.nonlinearity(self.rbr_reparam(inputs)) if self.rbr_identity is None: id_out = 0 else: id_out = self.rbr_identity(inputs) return self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out) def _fuse_bn(self, branch): if branch is None: return 0, 0 if isinstance(branch, nn.Sequential): kernel = branch.conv.weight.detach().cpu().numpy() running_mean = branch.bn.running_mean.cpu().numpy() running_var = branch.bn.running_var.cpu().numpy() gamma = branch.bn.weight.detach().cpu().numpy() beta = branch.bn.bias.detach().cpu().numpy() eps = branch.bn.eps else: assert isinstance(branch, nn.BatchNorm2d) kernel = np.zeros((self.in_channels, self.in_channels, 3, 3)) for i in range(self.in_channels): kernel[i, i, 1, 1] = 1 running_mean = branch.running_mean.cpu().numpy() running_var = branch.running_var.cpu().numpy() gamma = branch.weight.detach().cpu().numpy() beta = branch.bias.detach().cpu().numpy() eps = branch.eps std = np.sqrt(running_var + eps) t = gamma / std t = np.reshape(t, (-1, 1, 1, 1)) t = np.tile(t, (1, kernel.shape[1], kernel.shape[2], kernel.shape[3])) return kernel * t, beta - running_mean * gamma / std def _pad_1x1_to_3x3(self, kernel1x1): if kernel1x1 is None: return 0 kernel = np.zeros((kernel1x1.shape[0], kernel1x1.shape[1], 3, 3)) kernel[:, :, 1:2, 1:2] = kernel1x1 return kernel def repvgg_convert(self): kernel3x3, bias3x3 = self._fuse_bn(self.rbr_dense) kernel1x1, bias1x1 = self._fuse_bn(self.rbr_1x1) kernelid, biasid = self._fuse_bn(self.rbr_identity) return kernel3x3 + self._pad_1x1_to_3x3(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid class RepVGG(nn.Module): def __init__(self, num_blocks, num_classes=1000, width_multiplier=None, override_groups_map=None, deploy=False): super(RepVGG, self).__init__() assert len(width_multiplier) == 4 self.deploy = deploy self.override_groups_map = override_groups_map or dict() assert 0 not in self.override_groups_map self.in_planes = min(64, int(64 * width_multiplier[0])) self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy) self.cur_layer_idx = 1 self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2) self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2) self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2) self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2) self.gap = nn.AvgPool2d(7) self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes) def _make_stage(self, planes, num_blocks, stride): strides = [stride] + [1]*(num_blocks-1) blocks = [] for stride in strides: cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1) blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3, stride=stride, padding=1, groups=cur_groups, deploy=self.deploy)) self.in_planes = planes self.cur_layer_idx += 1 return nn.Sequential(*blocks) def forward(self, x): out = self.stage0(x) out = self.stage1(out) out = self.stage2(out) out = self.stage3(out) out = self.stage4(out) out = self.gap(out) out = out.view(out.size(0), -1) out = self.linear(out) return out optional_groupwise_layers = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26] g2_map = {l: 2 for l in optional_groupwise_layers} g4_map = {l: 4 for l in optional_groupwise_layers} def create_RepVGG_A0(deploy=False): return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000, width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy) def create_RepVGG_A1(deploy=False): return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000, width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy) def create_RepVGG_A2(deploy=False): return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000, width_multiplier=[1.5, 1.5, 1.5, 2.75], override_groups_map=None, deploy=deploy) def create_RepVGG_B0(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy) def create_RepVGG_B1(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2, 2, 2, 4], override_groups_map=None, deploy=deploy) def create_RepVGG_B1g2(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2, 2, 2, 4], override_groups_map=g2_map, deploy=deploy) def create_RepVGG_B1g4(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2, 2, 2, 4], override_groups_map=g4_map, deploy=deploy) def create_RepVGG_B2(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy) def create_RepVGG_B2g2(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g2_map, deploy=deploy) def create_RepVGG_B2g4(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g4_map, deploy=deploy) def create_RepVGG_B3(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[3, 3, 3, 5], override_groups_map=None, deploy=deploy) def create_RepVGG_B3g2(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[3, 3, 3, 5], override_groups_map=g2_map, deploy=deploy) def create_RepVGG_B3g4(deploy=False): return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000, width_multiplier=[3, 3, 3, 5], override_groups_map=g4_map, deploy=deploy) # Use like this: # train_model = create_RepVGG_A0(deploy=False) # train train_model # deploy_model = repvgg_convert(train_model, create_RepVGG_A0, save_path='repvgg_deploy.pth') def repvgg_convert(model:torch.nn.Module, build_func, save_path=None): converted_weights = {} for name, module in model.named_modules(): if hasattr(module, 'repvgg_convert'): kernel, bias = module.repvgg_convert() converted_weights[name + '.rbr_reparam.weight'] = kernel converted_weights[name + '.rbr_reparam.bias'] = bias elif isinstance(module, torch.nn.Linear): converted_weights[name + '.weight'] = module.weight.detach().cpu().numpy() converted_weights[name + '.bias'] = module.bias.detach().cpu().numpy() else: print(name, type(module)) del model deploy_model = build_func(deploy=True) for name, param in deploy_model.named_parameters(): print('deploy param: ', name, param.size(), np.mean(converted_weights[name])) param.data = torch.from_numpy(converted_weights[name]).float() if save_path is not None and save_path.endswith('pth'): torch.save(deploy_model.state_dict(), save_path) return deploy_model
4. 实验结果
4.1 ImageNet分类
4.2 语义分割
5 参考
[1].RepVGG: Making VGG-style ConvNets Great Again.
[2].https://github.com/DingXiaoH/RepVGG.
