diff --git a/lib/models/ImageNet_ResNet.py b/lib/models/ImageNet_ResNet.py new file mode 100644 index 0000000..9042db5 --- /dev/null +++ b/lib/models/ImageNet_ResNet.py @@ -0,0 +1,172 @@ +# Deep Residual Learning for Image Recognition, CVPR 2016 +import torch.nn as nn +from .initialization import initialize_resnet + +def conv3x3(in_planes, out_planes, stride=1, groups=1): + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, groups=groups, bias=False) + + +def conv1x1(in_planes, out_planes, stride=1): + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64): + super(BasicBlock, self).__init__() + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes) + self.bn2 = nn.BatchNorm2d(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = 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 += identity + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64): + super(Bottleneck, self).__init__() + width = int(planes * (base_width / 64.)) * groups + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv1x1(inplanes, width) + self.bn1 = nn.BatchNorm2d(width) + self.conv2 = conv3x3(width, width, stride, groups) + self.bn2 = nn.BatchNorm2d(width) + self.conv3 = conv1x1(width, planes * self.expansion) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + + out = 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 += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + + def __init__(self, block_name, layers, deep_stem, num_classes, zero_init_residual, groups, width_per_group): + super(ResNet, self).__init__() + + #planes = [int(width_per_group * groups * 2 ** i) for i in range(4)] + if block_name == 'BasicBlock' : block= BasicBlock + elif block_name == 'Bottleneck': block= Bottleneck + else : raise ValueError('invalid block-name : {:}'.format(block_name)) + + if not deep_stem: + self.conv = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), + nn.BatchNorm2d(64), nn.ReLU(inplace=True)) + else: + self.conv = nn.Sequential( + nn.Conv2d( 3, 32, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(32), nn.ReLU(inplace=True), + nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(32), nn.ReLU(inplace=True), + nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1, bias=False), + nn.BatchNorm2d(64), nn.ReLU(inplace=True)) + self.inplanes = 64 + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64 , layers[0], stride=1, groups=groups, base_width=width_per_group) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, groups=groups, base_width=width_per_group) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, groups=groups, base_width=width_per_group) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, groups=groups, base_width=width_per_group) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + self.message = 'block = {:}, layers = {:}, deep_stem = {:}, num_classes = {:}'.format(block, layers, deep_stem, num_classes) + + self.apply( initialize_resnet ) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck): + nn.init.constant_(m.bn3.weight, 0) + elif isinstance(m, BasicBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride, groups, base_width): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + if stride == 2: + downsample = nn.Sequential( + nn.AvgPool2d(kernel_size=2, stride=2, padding=0), + conv1x1(self.inplanes, planes * block.expansion, 1), + nn.BatchNorm2d(planes * block.expansion), + ) + elif stride == 1: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + nn.BatchNorm2d(planes * block.expansion), + ) + else: raise ValueError('invalid stride [{:}] for downsample'.format(stride)) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample, groups, base_width)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append(block(self.inplanes, planes, 1, None, groups, base_width)) + + return nn.Sequential(*layers) + + def get_message(self): + return self.message + + def forward(self, x): + x = self.conv(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + features = self.avgpool(x) + features = features.view(features.size(0), -1) + logits = self.fc(features) + + return features, logits diff --git a/lib/models/__init__.py b/lib/models/__init__.py index 4f7b735..b17934d 100644 --- a/lib/models/__init__.py +++ b/lib/models/__init__.py @@ -109,7 +109,7 @@ def get_cifar_models(config, extra_path=None): def get_imagenet_models(config): super_type = getattr(config, 'super_type', 'basic') if super_type == 'basic': - from .ImagenetResNet import ResNet + from .ImageNet_ResNet import ResNet from .ImageNet_MobileNetV2 import MobileNetV2 if config.arch == 'resnet': return ResNet(config.block_name, config.layers, config.deep_stem, config.class_num, config.zero_init_residual, config.groups, config.width_per_group)