Source code for models.slca_utils.convs.cifar_resnet

'''
Reference:
https://github.com/khurramjaved96/incremental-learning/blob/autoencoders/model/resnet32.py
'''
import math

import torch
import torch.nn as nn
import torch.nn.functional as F


[docs] class DownsampleA(nn.Module): def __init__(self, nIn, nOut, stride): super(DownsampleA, self).__init__() assert stride == 2 self.avg = nn.AvgPool2d(kernel_size=1, stride=stride)
[docs] def forward(self, x): x = self.avg(x) return torch.cat((x, x.mul(0)), 1)
[docs] class DownsampleB(nn.Module): def __init__(self, nIn, nOut, stride): super(DownsampleB, self).__init__() self.conv = nn.Conv2d(nIn, nOut, kernel_size=1, stride=stride, padding=0, bias=False) self.bn = nn.BatchNorm2d(nOut)
[docs] def forward(self, x): x = self.conv(x) x = self.bn(x) return x
[docs] class DownsampleC(nn.Module): def __init__(self, nIn, nOut, stride): super(DownsampleC, self).__init__() assert stride != 1 or nIn != nOut self.conv = nn.Conv2d(nIn, nOut, kernel_size=1, stride=stride, padding=0, bias=False)
[docs] def forward(self, x): x = self.conv(x) return x
[docs] class DownsampleD(nn.Module): def __init__(self, nIn, nOut, stride): super(DownsampleD, self).__init__() assert stride == 2 self.conv = nn.Conv2d(nIn, nOut, kernel_size=2, stride=stride, padding=0, bias=False) self.bn = nn.BatchNorm2d(nOut)
[docs] def forward(self, x): x = self.conv(x) x = self.bn(x) return x
[docs] class ResNetBasicblock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(ResNetBasicblock, self).__init__() self.conv_a = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn_a = nn.BatchNorm2d(planes) self.conv_b = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) self.bn_b = nn.BatchNorm2d(planes) self.downsample = downsample
[docs] def forward(self, x): residual = x basicblock = self.conv_a(x) basicblock = self.bn_a(basicblock) basicblock = F.relu(basicblock, inplace=True) basicblock = self.conv_b(basicblock) basicblock = self.bn_b(basicblock) if self.downsample is not None: residual = self.downsample(x) return F.relu(residual + basicblock, inplace=True)
[docs] class CifarResNet(nn.Module): """ ResNet optimized for the Cifar Dataset, as specified in https://arxiv.org/abs/1512.03385.pdf """ def __init__(self, block, depth, channels=3): super(CifarResNet, self).__init__() # Model type specifies number of layers for CIFAR-10 and CIFAR-100 model assert (depth - 2) % 6 == 0, 'depth should be one of 20, 32, 44, 56, 110' layer_blocks = (depth - 2) // 6 self.conv_1_3x3 = nn.Conv2d(channels, 16, kernel_size=3, stride=1, padding=1, bias=False) self.bn_1 = nn.BatchNorm2d(16) self.inplanes = 16 self.stage_1 = self._make_layer(block, 16, layer_blocks, 1) self.stage_2 = self._make_layer(block, 32, layer_blocks, 2) self.stage_3 = self._make_layer(block, 64, layer_blocks, 2) self.avgpool = nn.AvgPool2d(8) self.out_dim = 64 * block.expansion self.fc = nn.Linear(64 * block.expansion, 10) for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) # m.bias.data.zero_() elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() elif isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight) m.bias.data.zero_() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = DownsampleA(self.inplanes, planes * block.expansion, stride) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers)
[docs] def forward(self, x): x = self.conv_1_3x3(x) # [bs, 16, 32, 32] x = F.relu(self.bn_1(x), inplace=True) x_1 = self.stage_1(x) # [bs, 16, 32, 32] x_2 = self.stage_2(x_1) # [bs, 32, 16, 16] x_3 = self.stage_3(x_2) # [bs, 64, 8, 8] pooled = self.avgpool(x_3) # [bs, 64, 1, 1] features = pooled.view(pooled.size(0), -1) # [bs, 64] return { 'fmaps': [x_1, x_2, x_3], 'features': features }
@property def last_conv(self): return self.stage_3[-1].conv_b
[docs] def resnet20mnist(): """Constructs a ResNet-20 model for MNIST.""" model = CifarResNet(ResNetBasicblock, 20, 1) return model
[docs] def resnet32mnist(): """Constructs a ResNet-32 model for MNIST.""" model = CifarResNet(ResNetBasicblock, 32, 1) return model
[docs] def resnet20(): """Constructs a ResNet-20 model for CIFAR-10.""" model = CifarResNet(ResNetBasicblock, 20) return model
[docs] def resnet32(): """Constructs a ResNet-32 model for CIFAR-10.""" model = CifarResNet(ResNetBasicblock, 32) return model
[docs] def resnet44(): """Constructs a ResNet-44 model for CIFAR-10.""" model = CifarResNet(ResNetBasicblock, 44) return model
[docs] def resnet56(): """Constructs a ResNet-56 model for CIFAR-10.""" model = CifarResNet(ResNetBasicblock, 56) return model
[docs] def resnet110(): """Constructs a ResNet-110 model for CIFAR-10.""" model = CifarResNet(ResNetBasicblock, 110) return model