diff --git a/imagenet.py b/imagenet.py index f0c5c0a..edede9e 100644 --- a/imagenet.py +++ b/imagenet.py @@ -329,11 +329,11 @@ def test(val_loader, model, criterion, epoch, use_cuda): bar.finish() return (losses.avg, top1.avg, top5.avg) -def save_checkpoint(state, is_best, checkpoint='checkpoint', filename='checkpoint.pth.tar'): +def save_checkpoint(state, is_best, checkpoint='checkpoint', filename='checkpoint.pth'): filepath = os.path.join(checkpoint, filename) torch.save(state, filepath) if is_best: - shutil.copyfile(filepath, os.path.join(checkpoint, 'model_best.pth.tar')) + shutil.copyfile(filepath, os.path.join(checkpoint, 'model_best.pth')) def adjust_learning_rate(optimizer, epoch): global state diff --git a/models/imagenet/__init__.py b/models/imagenet/__init__.py index 9ad77bf..cffa6e7 100644 --- a/models/imagenet/__init__.py +++ b/models/imagenet/__init__.py @@ -9,3 +9,4 @@ from .resnext_ibn_a import * from .se_resnet import * from .se_resnet_ibn_a import * +from .se_resnet_ibn_b import * diff --git a/models/imagenet/se_resnet_ibn_b.py b/models/imagenet/se_resnet_ibn_b.py new file mode 100644 index 0000000..7bcf619 --- /dev/null +++ b/models/imagenet/se_resnet_ibn_b.py @@ -0,0 +1,190 @@ +from .se_module import SELayer +import torch.nn as nn +import torch +import math + +__all__ = ['se_resnet50_ibn_b', 'se_resnet101_ibn_b', 'se_resnet152_ibn_b'] + +def conv3x3(in_planes, out_planes, stride=1): + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) + + +class SEBasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16, IN=False): + super(SEBasicBlock, self).__init__() + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = nn.BatchNorm2d(planes) + self.relu = nn.ReLU(inplace=True) + self.conv2 = conv3x3(planes, planes, 1) + self.bn2 = nn.BatchNorm2d(planes) + self.IN = None + if IN: + self.IN = nn.InstanceNorm2d(planes * 4, affine=True) + self.se = SELayer(planes, reduction) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.se(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + if self.IN is not None: + out = self.IN(out) + out = self.relu(out) + + return out + + +class SEBottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16, IN=False): + super(SEBottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.IN = None + if IN: + self.IN = nn.InstanceNorm2d(planes * 4, affine=True) + self.relu = nn.ReLU(inplace=True) + self.se = SELayer(planes * 4, reduction) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = 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) + out = self.se(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + if self.IN is not None: + out = self.IN(out) + out = self.relu(out) + + return out + +class ResNet(nn.Module): + + def __init__(self, block, layers, num_classes=1000): + self.inplanes = 64 + super(ResNet, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.InstanceNorm2d(64, affine=True) + 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], stride=1, IN=True) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, IN=True) + 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.AvgPool2d(7) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + self.conv1.weight.data.normal_(0, math.sqrt(2. / (7 * 7 * 64))) + 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)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.InstanceNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1, IN=False): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks-1): + layers.append(block(self.inplanes, planes)) + layers.append(block(self.inplanes, planes, IN=IN)) + + 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 x + + +def se_resnet50_ibn_b(num_classes=1000): + """Constructs a ResNet-50 model. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet(SEBottleneck, [3, 4, 6, 3], num_classes=num_classes) + model.avgpool = nn.AdaptiveAvgPool2d(1) + return model + + +def se_resnet101_ibn_b(num_classes=1000): + """Constructs a ResNet-101 model. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet(SEBottleneck, [3, 4, 23, 3], num_classes=num_classes) + model.avgpool = nn.AdaptiveAvgPool2d(1) + return model + + +def se_resnet152_ibn_b(num_classes): + """Constructs a ResNet-152 model. + + Args: + pretrained (bool): If True, returns a model pre-trained on ImageNet + """ + model = ResNet(SEBottleneck, [3, 8, 36, 3], num_classes=num_classes) + model.avgpool = nn.AdaptiveAvgPool2d(1) + return model diff --git a/test.sh b/test.sh index 7f43b35..052feb8 100755 --- a/test.sh +++ b/test.sh @@ -5,7 +5,7 @@ data_path=/pathToYourImageNetDataset/ python -u imagenet.py \ -a $model \ --test-batch 100 \ - --model_weight pretrained/${model}.pth.tar \ + --model_weight pretrained/${model}.pth \ -e \ -j 16 \ --data $data_path \