| from collections import namedtuple |
| import torch |
| import torch.nn as nn |
| from torch.nn import Dropout |
| from torch.nn import MaxPool2d |
| from torch.nn import Sequential |
| from torch.nn import Conv2d, Linear |
| from torch.nn import BatchNorm1d, BatchNorm2d |
| from torch.nn import ReLU, Sigmoid |
| from torch.nn import Module |
| from torch.nn import PReLU |
| import os |
|
|
| def build_model(model_name='ir_50'): |
| if model_name == 'ir_101': |
| return IR_101(input_size=(112,112)) |
| elif model_name == 'ir_50': |
| return IR_50(input_size=(112,112)) |
| elif model_name == 'ir_se_50': |
| return IR_SE_50(input_size=(112,112)) |
| elif model_name == 'ir_34': |
| return IR_34(input_size=(112,112)) |
| elif model_name == 'ir_18': |
| return IR_18(input_size=(112,112)) |
| else: |
| raise ValueError('not a correct model name', model_name) |
|
|
| def initialize_weights(modules): |
| """ Weight initilize, conv2d and linear is initialized with kaiming_normal |
| """ |
| for m in modules: |
| if isinstance(m, nn.Conv2d): |
| nn.init.kaiming_normal_(m.weight, |
| mode='fan_out', |
| nonlinearity='relu') |
| if m.bias is not None: |
| 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, |
| mode='fan_out', |
| nonlinearity='relu') |
| if m.bias is not None: |
| m.bias.data.zero_() |
|
|
|
|
| class Flatten(Module): |
| """ Flat tensor |
| """ |
| def forward(self, input): |
| return input.view(input.size(0), -1) |
|
|
|
|
| class LinearBlock(Module): |
| """ Convolution block without no-linear activation layer |
| """ |
| def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1): |
| super(LinearBlock, self).__init__() |
| self.conv = Conv2d(in_c, out_c, kernel, stride, padding, groups=groups, bias=False) |
| self.bn = BatchNorm2d(out_c) |
|
|
| def forward(self, x): |
| x = self.conv(x) |
| x = self.bn(x) |
| return x |
|
|
|
|
| class GNAP(Module): |
| """ Global Norm-Aware Pooling block |
| """ |
| def __init__(self, in_c): |
| super(GNAP, self).__init__() |
| self.bn1 = BatchNorm2d(in_c, affine=False) |
| self.pool = nn.AdaptiveAvgPool2d((1, 1)) |
| self.bn2 = BatchNorm1d(in_c, affine=False) |
|
|
| def forward(self, x): |
| x = self.bn1(x) |
| x_norm = torch.norm(x, 2, 1, True) |
| x_norm_mean = torch.mean(x_norm) |
| weight = x_norm_mean / x_norm |
| x = x * weight |
| x = self.pool(x) |
| x = x.view(x.shape[0], -1) |
| feature = self.bn2(x) |
| return feature |
|
|
|
|
| class GDC(Module): |
| """ Global Depthwise Convolution block |
| """ |
| def __init__(self, in_c, embedding_size): |
| super(GDC, self).__init__() |
| self.conv_6_dw = LinearBlock(in_c, in_c, |
| groups=in_c, |
| kernel=(7, 7), |
| stride=(1, 1), |
| padding=(0, 0)) |
| self.conv_6_flatten = Flatten() |
| self.linear = Linear(in_c, embedding_size, bias=False) |
| self.bn = BatchNorm1d(embedding_size, affine=False) |
|
|
| def forward(self, x): |
| x = self.conv_6_dw(x) |
| x = self.conv_6_flatten(x) |
| x = self.linear(x) |
| x = self.bn(x) |
| return x |
|
|
|
|
| class SEModule(Module): |
| """ SE block |
| """ |
| def __init__(self, channels, reduction): |
| super(SEModule, self).__init__() |
| self.avg_pool = nn.AdaptiveAvgPool2d(1) |
| self.fc1 = Conv2d(channels, channels // reduction, |
| kernel_size=1, padding=0, bias=False) |
|
|
| nn.init.xavier_uniform_(self.fc1.weight.data) |
|
|
| self.relu = ReLU(inplace=True) |
| self.fc2 = Conv2d(channels // reduction, channels, |
| kernel_size=1, padding=0, bias=False) |
|
|
| self.sigmoid = Sigmoid() |
|
|
| def forward(self, x): |
| module_input = x |
| x = self.avg_pool(x) |
| x = self.fc1(x) |
| x = self.relu(x) |
| x = self.fc2(x) |
| x = self.sigmoid(x) |
|
|
| return module_input * x |
|
|
|
|
|
|
| class BasicBlockIR(Module): |
| """ BasicBlock for IRNet |
| """ |
| def __init__(self, in_channel, depth, stride): |
| super(BasicBlockIR, self).__init__() |
| if in_channel == depth: |
| self.shortcut_layer = MaxPool2d(1, stride) |
| else: |
| self.shortcut_layer = Sequential( |
| Conv2d(in_channel, depth, (1, 1), stride, bias=False), |
| BatchNorm2d(depth)) |
| self.res_layer = Sequential( |
| BatchNorm2d(in_channel), |
| Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False), |
| BatchNorm2d(depth), |
| PReLU(depth), |
| Conv2d(depth, depth, (3, 3), stride, 1, bias=False), |
| BatchNorm2d(depth)) |
|
|
| def forward(self, x): |
| shortcut = self.shortcut_layer(x) |
| res = self.res_layer(x) |
|
|
| return res + shortcut |
|
|
|
|
| class BottleneckIR(Module): |
| """ BasicBlock with bottleneck for IRNet |
| """ |
| def __init__(self, in_channel, depth, stride): |
| super(BottleneckIR, self).__init__() |
| reduction_channel = depth // 4 |
| if in_channel == depth: |
| self.shortcut_layer = MaxPool2d(1, stride) |
| else: |
| self.shortcut_layer = Sequential( |
| Conv2d(in_channel, depth, (1, 1), stride, bias=False), |
| BatchNorm2d(depth)) |
| self.res_layer = Sequential( |
| BatchNorm2d(in_channel), |
| Conv2d(in_channel, reduction_channel, (1, 1), (1, 1), 0, bias=False), |
| BatchNorm2d(reduction_channel), |
| PReLU(reduction_channel), |
| Conv2d(reduction_channel, reduction_channel, (3, 3), (1, 1), 1, bias=False), |
| BatchNorm2d(reduction_channel), |
| PReLU(reduction_channel), |
| Conv2d(reduction_channel, depth, (1, 1), stride, 0, bias=False), |
| BatchNorm2d(depth)) |
|
|
| def forward(self, x): |
| shortcut = self.shortcut_layer(x) |
| res = self.res_layer(x) |
|
|
| return res + shortcut |
|
|
|
|
| class BasicBlockIRSE(BasicBlockIR): |
| def __init__(self, in_channel, depth, stride): |
| super(BasicBlockIRSE, self).__init__(in_channel, depth, stride) |
| self.res_layer.add_module("se_block", SEModule(depth, 16)) |
|
|
|
|
| class BottleneckIRSE(BottleneckIR): |
| def __init__(self, in_channel, depth, stride): |
| super(BottleneckIRSE, self).__init__(in_channel, depth, stride) |
| self.res_layer.add_module("se_block", SEModule(depth, 16)) |
|
|
|
|
| class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])): |
| '''A named tuple describing a ResNet block.''' |
|
|
|
|
| def get_block(in_channel, depth, num_units, stride=2): |
|
|
| return [Bottleneck(in_channel, depth, stride)] +\ |
| [Bottleneck(depth, depth, 1) for i in range(num_units - 1)] |
|
|
|
|
| def get_blocks(num_layers): |
| if num_layers == 18: |
| blocks = [ |
| get_block(in_channel=64, depth=64, num_units=2), |
| get_block(in_channel=64, depth=128, num_units=2), |
| get_block(in_channel=128, depth=256, num_units=2), |
| get_block(in_channel=256, depth=512, num_units=2) |
| ] |
| elif num_layers == 34: |
| blocks = [ |
| get_block(in_channel=64, depth=64, num_units=3), |
| get_block(in_channel=64, depth=128, num_units=4), |
| get_block(in_channel=128, depth=256, num_units=6), |
| get_block(in_channel=256, depth=512, num_units=3) |
| ] |
| elif num_layers == 50: |
| blocks = [ |
| get_block(in_channel=64, depth=64, num_units=3), |
| get_block(in_channel=64, depth=128, num_units=4), |
| get_block(in_channel=128, depth=256, num_units=14), |
| get_block(in_channel=256, depth=512, num_units=3) |
| ] |
| elif num_layers == 100: |
| blocks = [ |
| get_block(in_channel=64, depth=64, num_units=3), |
| get_block(in_channel=64, depth=128, num_units=13), |
| get_block(in_channel=128, depth=256, num_units=30), |
| get_block(in_channel=256, depth=512, num_units=3) |
| ] |
| elif num_layers == 152: |
| blocks = [ |
| get_block(in_channel=64, depth=256, num_units=3), |
| get_block(in_channel=256, depth=512, num_units=8), |
| get_block(in_channel=512, depth=1024, num_units=36), |
| get_block(in_channel=1024, depth=2048, num_units=3) |
| ] |
| elif num_layers == 200: |
| blocks = [ |
| get_block(in_channel=64, depth=256, num_units=3), |
| get_block(in_channel=256, depth=512, num_units=24), |
| get_block(in_channel=512, depth=1024, num_units=36), |
| get_block(in_channel=1024, depth=2048, num_units=3) |
| ] |
|
|
| return blocks |
|
|
|
|
| class Backbone(Module): |
| def __init__(self, input_size, num_layers, mode='ir'): |
| """ Args: |
| input_size: input_size of backbone |
| num_layers: num_layers of backbone |
| mode: support ir or irse |
| """ |
| super(Backbone, self).__init__() |
| assert input_size[0] in [112, 224], \ |
| "input_size should be [112, 112] or [224, 224]" |
| assert num_layers in [18, 34, 50, 100, 152, 200], \ |
| "num_layers should be 18, 34, 50, 100 or 152" |
| assert mode in ['ir', 'ir_se'], \ |
| "mode should be ir or ir_se" |
| self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False), |
| BatchNorm2d(64), PReLU(64)) |
| blocks = get_blocks(num_layers) |
| if num_layers <= 100: |
| if mode == 'ir': |
| unit_module = BasicBlockIR |
| elif mode == 'ir_se': |
| unit_module = BasicBlockIRSE |
| output_channel = 512 |
| else: |
| if mode == 'ir': |
| unit_module = BottleneckIR |
| elif mode == 'ir_se': |
| unit_module = BottleneckIRSE |
| output_channel = 2048 |
|
|
| if input_size[0] == 112: |
| self.output_layer = Sequential(BatchNorm2d(output_channel), |
| Dropout(0.4), Flatten(), |
| Linear(output_channel * 7 * 7, 512), |
| BatchNorm1d(512, affine=False)) |
| else: |
| self.output_layer = Sequential( |
| BatchNorm2d(output_channel), Dropout(0.4), Flatten(), |
| Linear(output_channel * 14 * 14, 512), |
| BatchNorm1d(512, affine=False)) |
|
|
| modules = [] |
| for block in blocks: |
| for bottleneck in block: |
| modules.append( |
| unit_module(bottleneck.in_channel, bottleneck.depth, |
| bottleneck.stride)) |
| self.body = Sequential(*modules) |
|
|
| initialize_weights(self.modules()) |
|
|
|
|
| def forward(self, x): |
| |
| |
| |
| x = self.input_layer(x) |
|
|
| for idx, module in enumerate(self.body): |
| x = module(x) |
|
|
| x = self.output_layer(x) |
| norm = torch.norm(x, 2, 1, True) |
| output = torch.div(x, norm) |
|
|
| return output, norm |
|
|
|
|
|
|
| def IR_18(input_size): |
| """ Constructs a ir-18 model. |
| """ |
| model = Backbone(input_size, 18, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_34(input_size): |
| """ Constructs a ir-34 model. |
| """ |
| model = Backbone(input_size, 34, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_50(input_size): |
| """ Constructs a ir-50 model. |
| """ |
| model = Backbone(input_size, 50, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_101(input_size): |
| """ Constructs a ir-101 model. |
| """ |
| model = Backbone(input_size, 100, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_152(input_size): |
| """ Constructs a ir-152 model. |
| """ |
| model = Backbone(input_size, 152, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_200(input_size): |
| """ Constructs a ir-200 model. |
| """ |
| model = Backbone(input_size, 200, 'ir') |
|
|
| return model |
|
|
|
|
| def IR_SE_50(input_size): |
| """ Constructs a ir_se-50 model. |
| """ |
| model = Backbone(input_size, 50, 'ir_se') |
|
|
| return model |
|
|
|
|
| def IR_SE_101(input_size): |
| """ Constructs a ir_se-101 model. |
| """ |
| model = Backbone(input_size, 100, 'ir_se') |
|
|
| return model |
|
|
|
|
| def IR_SE_152(input_size): |
| """ Constructs a ir_se-152 model. |
| """ |
| model = Backbone(input_size, 152, 'ir_se') |
|
|
| return model |
|
|
|
|
| def IR_SE_200(input_size): |
| """ Constructs a ir_se-200 model. |
| """ |
| model = Backbone(input_size, 200, 'ir_se') |
|
|
| return model |
|
|
|
|
