Towards Stabilizing Batch Statistics in Backward Propagation of Batch Normalization
MABN是一种新颖的归一化方法，性能优于SynCBN、BRN、BN等方法。作者团队旷视科技&复旦，该方法一小时前开源！
点击获取论文，github链接。
这里只简单介绍MABN和它的使用。
论文摘要
BN是深度学习领域中使用最广泛的技术之一。但是由于批次大小不足，其性能可能会严重下降。这种弱点限制了BN在许多计算机视觉任务上的使用，由于内存消耗的限制，批处理的大小通常很小。因此，已经提出了许多改进的归一化技术，它们要么不能完全恢复BN的性能，要么必须在推理过程中引入额外的非线性运算并增加巨大的消耗。
本文中，我们揭示了BN向后传播涉及两个额外的批处理统计信息。与梯度相关的额外批处理统计信息会严重影响深度神经网络的训练，根据分析，我们提出了一种新颖的归一化方法，称之为移动平均批量归一化（MABN）。在小批量情况下，MABN可以完全恢复vanilla BN的性能，而无需在推理过程中引入任何其它非线性操作。我们通过理论分析和实验证明了MABN的好处。通过实验，我们证明了MABN在包括Imagenet和COCO在内的多种计算机视觉任务的有效性。
实验结果
由上图可以看出，在Batch Size很小的情况下，MABN也得到了很好的效果。
MABN在batch_size很小等于2的时候，对于Imagenet数据集，相比BN和BRN，得到的Top 1 Accuracy与使用BN且batch_size=32时相当。
对于COCO数据集，MABN仅在batch_size=2就可获得了其它归一化方法需要将batch_size设置为16才能获得的精度，进一步说明了MABN在小批量数据下的有效性和优良性。
MABN的简单使用
以resnet50为例，将MABN替代resnet50网络构建用到的nn.BatchNorm2d。
将MABN.py文件包括在工作目录，
重点注意from networks.MABN import MABN2d as BatchNorm2d
resnet.py改写为：
import torch
import torch.nn as nn
import torch.nn.functional as Ffrom networks.MABN import MABN2d as BatchNorm2dclass Conv2d(nn.Module):"""Conv2d layer with Weight Centralization"""def __init__(self, in_planes, out_planes, kernel_size=3, stride=1,padding=0, dilation=1, groups=1, bias=False):super(Conv2d, self).__init__()self.in_planes = in_planesself.out_planes = out_planesself.stride = strideself.padding = paddingself.dilation = dilationself.groups = groupsself.weight = nn.Parameter(torch.randn(out_planes, in_planes//groups, kernel_size, kernel_size))if bias:self.bias = nn.Parameter(torch.randn(out_planes))else:self.register_parameter('bias', None)def forward(self, x):weight = self.weightweight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True)weight = weight - weight_meanreturn F.conv2d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)def conv3x3(in_planes, out_planes, stride=1):"""3x3 convolution with padding"""return Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,padding=1, bias=False)def conv1x1(in_planes, out_planes, stride=1):"""1x1 convolution"""return Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)class BasicBlock(nn.Module):expansion = 1def __init__(self, inplanes, planes, stride=1, downsample=None):super(BasicBlock, self).__init__()self.conv1 = conv3x3(inplanes, planes, stride)self.bn1 = BatchNorm2d(planes)self.relu = nn.ReLU(inplace=True)self.conv2 = conv3x3(planes, planes)self.bn2 = BatchNorm2d(planes)self.downsample = downsampleself.stride = stridedef forward(self, x):identity = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)if self.downsample is not None:identity = self.downsample(x)out += identityout = self.relu(out)return outclass Bottleneck(nn.Module):expansion = 4def __init__(self, inplanes, planes, stride=1, downsample=None):super(Bottleneck, self).__init__()self.conv1 = conv1x1(inplanes, planes)self.bn1 = BatchNorm2d(planes)self.conv2 = conv3x3(planes, planes, stride)self.bn2 = BatchNorm2d(planes)self.conv3 = conv1x1(planes, planes * self.expansion)self.bn3 = BatchNorm2d(planes * self.expansion)self.relu = nn.ReLU(inplace=True)self.downsample = downsampleself.stride = stridedef forward(self, x):identity = xout = self.conv1(x)out = self.bn1(out)out = self.relu(out)out = self.conv2(out)out = self.bn2(out)out = self.relu(out)out = self.conv3(out)out = self.bn3(out)if self.downsample is not None:identity = self.downsample(x)out += identityout = self.relu(out)return outclass ResNet(nn.Module):def __init__(self, block, layers, num_classes=1000):super(ResNet, self).__init__()self.inplanes = 64self.conv1 = Conv2d(3, 64, kernel_size=7, stride=2, padding=3,bias=False)self.bn1 = BatchNorm2d(64)self.relu = nn.ReLU(inplace=True)self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)self.layer1 = self._make_layer(block, 64, layers[0])self.layer2 = self._make_layer(block, 128, layers[1], stride=2)self.layer3 = self._make_layer(block, 256, layers[2], stride=2)self.layer4 = self._make_layer(block, 512, layers[3], stride=2)self.avgpool = nn.AdaptiveAvgPool2d((1, 1))self.fc = nn.Linear(512 * block.expansion, num_classes)for m in self.modules():if isinstance(m, Conv2d):nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')elif isinstance(m, BatchNorm2d):nn.init.constant_(m.weight, 1)nn.init.constant_(m.bias, 0)def _make_layer(self, block, planes, blocks, stride=1):downsample = Noneif stride != 1 or self.inplanes != planes * block.expansion:downsample = nn.Sequential(conv1x1(self.inplanes, planes * block.expansion, stride),BatchNorm2d(planes * block.expansion),)layers = []layers.append(block(self.inplanes, planes, stride, downsample))self.inplanes = planes * block.expansionfor _ in range(1, blocks):layers.append(block(self.inplanes, planes))return nn.Sequential(*layers)def forward(self, x):x = self.conv1(x)x = self.bn1(x)x = self.relu(x)x = self.maxpool(x)x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)x = self.layer4(x)x = self.avgpool(x)x = x.view(x.size(0), -1)x = self.fc(x)return xdef resnet50(**kwargs):"""Constructs a ResNet-50 model.Args:pretrained (bool): If True, returns a model pre-trained on ImageNet"""model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)return model
那么返回的resnet50就是用了MABN的resnet50。