# The implementation is adopted from face.evoLVe, made publicly available under the MIT License # at https://github.com/ZhaoJ9014/face.evoLVe/blob/master/backbone/model_irse.py from collections import namedtuple import torch import torch.nn as nn from torch.nn import (AdaptiveAvgPool2d, BatchNorm1d, BatchNorm2d, Conv2d, Dropout, Linear, MaxPool2d, Module, PReLU, ReLU, Sequential, Sigmoid) class Flatten(Module): def forward(self, input): return input.view(input.size(0), -1) def l2_norm(input, axis=1): norm = torch.norm(input, 2, axis, True) output = torch.div(input, norm) return output class SEModule(Module): def __init__(self, channels, reduction): super(SEModule, self).__init__() self.avg_pool = 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 bottleneck_IR(Module): def __init__(self, in_channel, depth, stride): super(bottleneck_IR, 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), 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 bottleneck_IR_SE(Module): def __init__(self, in_channel, depth, stride): super(bottleneck_IR_SE, 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), PReLU(depth), Conv2d(depth, depth, (3, 3), stride, 1, bias=False), BatchNorm2d(depth), SEModule(depth, 16)) def forward(self, x): shortcut = self.shortcut_layer(x) res = self.res_layer(x) return res + shortcut 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 == 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=64, num_units=3), get_block(in_channel=64, depth=128, num_units=8), get_block(in_channel=128, depth=256, num_units=36), get_block(in_channel=256, depth=512, num_units=3) ] return blocks class Backbone(Module): def __init__(self, input_size, num_layers, mode='ir'): super(Backbone, self).__init__() assert input_size[0] in [ 112, 224 ], 'input_size should be [112, 112] or [224, 224]' assert num_layers in [50, 100, 152], 'num_layers should be 50, 100 or 152' assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se' blocks = get_blocks(num_layers) if mode == 'ir': unit_module = bottleneck_IR elif mode == 'ir_se': unit_module = bottleneck_IR_SE self.input_layer = Sequential( Conv2d(3, 64, (3, 3), 1, 1, bias=False), BatchNorm2d(64), PReLU(64)) if input_size[0] == 112: self.output_layer = Sequential( BatchNorm2d(512), Dropout(0.4), Flatten(), Linear(512 * 7 * 7, 512), BatchNorm1d(512, affine=False)) else: self.output_layer = Sequential( BatchNorm2d(512), Dropout(0.4), Flatten(), Linear(512 * 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) self._initialize_weights() def forward(self, x): x = self.input_layer(x) x = self.body(x) conv_out = x.view(x.shape[0], -1) x = self.output_layer(x) return l2_norm(x), conv_out def _initialize_weights(self): for m in self.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_() 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_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