From aebb6033e0ed7f27e0fa0a44e6533a134a070f1a Mon Sep 17 00:00:00 2001 From: KaiyangZhou Date: Sun, 10 Nov 2019 21:31:28 +0000 Subject: [PATCH] add ibn-net --- torchreid/models/__init__.py | 4 + torchreid/models/resnet_ibn_a.py | 251 +++++++++++++++++++++++++++++++ torchreid/models/resnet_ibn_b.py | 233 ++++++++++++++++++++++++++++ 3 files changed, 488 insertions(+) create mode 100644 torchreid/models/resnet_ibn_a.py create mode 100644 torchreid/models/resnet_ibn_b.py diff --git a/torchreid/models/__init__.py b/torchreid/models/__init__.py index 67a285e..347653b 100644 --- a/torchreid/models/__init__.py +++ b/torchreid/models/__init__.py @@ -4,6 +4,8 @@ import torch from .resnet import * from .resnetmid import * +from .resnet_ibn_a import * +from .resnet_ibn_b import * from .senet import * from .densenet import * from .inceptionresnetv2 import * @@ -47,6 +49,8 @@ __model_factory = { 'inceptionresnetv2': inceptionresnetv2, 'inceptionv4': inceptionv4, 'xception': xception, + 'resnet50_ibn_a': resnet50_ibn_a, + 'resnet50_ibn_b': resnet50_ibn_b, # lightweight models 'nasnsetmobile': nasnetamobile, 'mobilenetv2_x1_0': mobilenetv2_x1_0, diff --git a/torchreid/models/resnet_ibn_a.py b/torchreid/models/resnet_ibn_a.py new file mode 100644 index 0000000..0b72dbd --- /dev/null +++ b/torchreid/models/resnet_ibn_a.py @@ -0,0 +1,251 @@ +""" +Credit to https://github.com/XingangPan/IBN-Net. +""" +from __future__ import absolute_import +from __future__ import division + +__all__ = ['resnet50_ibn_a'] + +import torch +import torch.nn as nn +import math +import torch.utils.model_zoo as model_zoo + + +model_urls = { + 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', + 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', + 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', +} + + +def conv3x3(in_planes, out_planes, stride=1): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + 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): + residual = 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: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class IBN(nn.Module): + def __init__(self, planes): + super(IBN, self).__init__() + half1 = int(planes/2) + self.half = half1 + half2 = planes - half1 + self.IN = nn.InstanceNorm2d(half1, affine=True) + self.BN = nn.BatchNorm2d(half2) + + def forward(self, x): + split = torch.split(x, self.half, 1) + out1 = self.IN(split[0].contiguous()) + out2 = self.BN(split[1].contiguous()) + out = torch.cat((out1, out2), 1) + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + if ibn: + self.bn1 = IBN(planes) + else: + 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 * self.expansion, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + 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) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + """Residual network + IBN layer. + + Reference: + - He et al. Deep Residual Learning for Image Recognition. CVPR 2016. + - Pan et al. Two at Once: Enhancing Learning and Generalization + Capacities via IBN-Net. ECCV 2018. + """ + + def __init__(self, block, layers, num_classes=1000, loss='softmax', + fc_dims=None, dropout_p=None, **kwargs): + scale = 64 + self.inplanes = scale + super(ResNet, self).__init__() + self.loss = loss + self.feature_dim = scale * 8 * block.expansion + + self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(scale) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, scale, layers[0]) + self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2) + self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) + self.layer4 = self._make_layer(block, scale*8, layers[3], stride=2) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = self._construct_fc_layer(fc_dims, scale * 8 * block.expansion, dropout_p) + self.classifier = nn.Linear(self.feature_dim, num_classes) + + 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): + 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 = [] + ibn = True + if planes == 512: + ibn = False + layers.append(block(self.inplanes, planes, ibn, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, ibn)) + + return nn.Sequential(*layers) + + def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None): + """Constructs fully connected layer + + Args: + fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed + input_dim (int): input dimension + dropout_p (float): dropout probability, if None, dropout is unused + """ + if fc_dims is None: + self.feature_dim = input_dim + return None + + assert isinstance(fc_dims, (list, tuple)), 'fc_dims must be either list or tuple, but got {}'.format(type(fc_dims)) + + layers = [] + for dim in fc_dims: + layers.append(nn.Linear(input_dim, dim)) + layers.append(nn.BatchNorm1d(dim)) + layers.append(nn.ReLU(inplace=True)) + if dropout_p is not None: + layers.append(nn.Dropout(p=dropout_p)) + input_dim = dim + + self.feature_dim = fc_dims[-1] + + return nn.Sequential(*layers) + + def featuremaps(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) + return x + + def forward(self, x): + f = self.featuremaps(x) + v = self.avgpool(f) + v = v.view(v.size(0), -1) + if self.fc is not None: + v = self.fc(v) + if not self.training: + return v + y = self.classifier(v) + if self.loss == 'softmax': + return y + elif self.loss == 'triplet': + return y, v + else: + raise KeyError("Unsupported loss: {}".format(self.loss)) + + +def init_pretrained_weights(model, model_url): + """Initializes model with pretrained weights. + + Layers that don't match with pretrained layers in name or size are kept unchanged. + """ + pretrain_dict = model_zoo.load_url(model_url) + model_dict = model.state_dict() + pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} + model_dict.update(pretrain_dict) + model.load_state_dict(model_dict) + + +def resnet50_ibn_a(num_classes, loss='softmax', pretrained=False, **kwargs): + model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs) + if pretrained: + init_pretrained_weights(model, model_urls['resnet50']) + return model diff --git a/torchreid/models/resnet_ibn_b.py b/torchreid/models/resnet_ibn_b.py new file mode 100644 index 0000000..309d4bc --- /dev/null +++ b/torchreid/models/resnet_ibn_b.py @@ -0,0 +1,233 @@ +""" +Credit to https://github.com/XingangPan/IBN-Net. +""" +from __future__ import absolute_import +from __future__ import division + +__all__ = ['resnet50_ibn_b'] + +import torch.nn as nn +import math +import torch.utils.model_zoo as model_zoo + + +model_urls = { + 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', + 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', + 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', +} + + +def conv3x3(in_planes, out_planes, stride=1): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, + padding=1, bias=False) + + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(BasicBlock, self).__init__() + 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): + residual = 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: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class Bottleneck(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None, IN=False): + super(Bottleneck, 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 * self.expansion, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * self.expansion) + self.IN = None + if IN: + self.IN = nn.InstanceNorm2d(planes * 4, affine=True) + self.relu = nn.ReLU(inplace=True) + 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) + + 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): + """Residual network + IBN layer. + + Reference: + - He et al. Deep Residual Learning for Image Recognition. CVPR 2016. + - Pan et al. Two at Once: Enhancing Learning and Generalization + Capacities via IBN-Net. ECCV 2018. + """ + + def __init__(self, block, layers, num_classes=1000, loss='softmax', + fc_dims=None, dropout_p=None, **kwargs): + scale = 64 + self.inplanes = scale + super(ResNet, self).__init__() + self.loss = loss + self.feature_dim = scale * 8 * block.expansion + + self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.InstanceNorm2d(scale, 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, scale, layers[0], stride=1, IN=True) + self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2, IN=True) + self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2) + self.layer4 = self._make_layer(block, scale*8, layers[3], stride=2) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = self._construct_fc_layer(fc_dims, scale * 8 * block.expansion, dropout_p) + self.classifier = nn.Linear(self.feature_dim, num_classes) + + 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 _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None): + """Constructs fully connected layer + + Args: + fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed + input_dim (int): input dimension + dropout_p (float): dropout probability, if None, dropout is unused + """ + if fc_dims is None: + self.feature_dim = input_dim + return None + + assert isinstance(fc_dims, (list, tuple)), 'fc_dims must be either list or tuple, but got {}'.format(type(fc_dims)) + + layers = [] + for dim in fc_dims: + layers.append(nn.Linear(input_dim, dim)) + layers.append(nn.BatchNorm1d(dim)) + layers.append(nn.ReLU(inplace=True)) + if dropout_p is not None: + layers.append(nn.Dropout(p=dropout_p)) + input_dim = dim + + self.feature_dim = fc_dims[-1] + + return nn.Sequential(*layers) + + def featuremaps(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) + return x + + def forward(self, x): + f = self.featuremaps(x) + v = self.avgpool(f) + v = v.view(v.size(0), -1) + if self.fc is not None: + v = self.fc(v) + if not self.training: + return v + y = self.classifier(v) + if self.loss == 'softmax': + return y + elif self.loss == 'triplet': + return y, v + else: + raise KeyError("Unsupported loss: {}".format(self.loss)) + + +def init_pretrained_weights(model, model_url): + """Initializes model with pretrained weights. + + Layers that don't match with pretrained layers in name or size are kept unchanged. + """ + pretrain_dict = model_zoo.load_url(model_url) + model_dict = model.state_dict() + pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} + model_dict.update(pretrain_dict) + model.load_state_dict(model_dict) + + +def resnet50_ibn_b(num_classes, loss='softmax', pretrained=False, **kwargs): + model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs) + if pretrained: + init_pretrained_weights(model, model_urls['resnet50']) + return model