|
| 1 | +# encoding: utf-8 |
| 2 | +""" |
| 3 | +@author: xingyu liao |
| 4 | + |
| 5 | +""" |
| 6 | + |
| 7 | +import torch |
| 8 | +from torch import nn |
| 9 | + |
| 10 | +from fastreid.layers import get_norm |
| 11 | +from fastreid.modeling.backbones import BACKBONE_REGISTRY |
| 12 | + |
| 13 | + |
| 14 | +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): |
| 15 | + """3x3 convolution with padding""" |
| 16 | + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, |
| 17 | + padding=dilation, groups=groups, bias=False, dilation=dilation) |
| 18 | + |
| 19 | + |
| 20 | +def conv1x1(in_planes, out_planes, stride=1): |
| 21 | + """1x1 convolution""" |
| 22 | + return nn.Conv2d(in_planes, |
| 23 | + out_planes, |
| 24 | + kernel_size=1, |
| 25 | + stride=stride, |
| 26 | + bias=False) |
| 27 | + |
| 28 | + |
| 29 | +class IBasicBlock(nn.Module): |
| 30 | + expansion = 1 |
| 31 | + |
| 32 | + def __init__(self, inplanes, planes, bn_norm, stride=1, downsample=None, |
| 33 | + groups=1, base_width=64, dilation=1): |
| 34 | + super().__init__() |
| 35 | + if groups != 1 or base_width != 64: |
| 36 | + raise ValueError('BasicBlock only supports groups=1 and base_width=64') |
| 37 | + if dilation > 1: |
| 38 | + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") |
| 39 | + self.bn1 = get_norm(bn_norm, inplanes) |
| 40 | + self.conv1 = conv3x3(inplanes, planes) |
| 41 | + self.bn2 = get_norm(bn_norm, planes) |
| 42 | + self.prelu = nn.PReLU(planes) |
| 43 | + self.conv2 = conv3x3(planes, planes, stride) |
| 44 | + self.bn3 = get_norm(bn_norm, planes) |
| 45 | + self.downsample = downsample |
| 46 | + self.stride = stride |
| 47 | + |
| 48 | + def forward(self, x): |
| 49 | + identity = x |
| 50 | + out = self.bn1(x) |
| 51 | + out = self.conv1(out) |
| 52 | + out = self.bn2(out) |
| 53 | + out = self.prelu(out) |
| 54 | + out = self.conv2(out) |
| 55 | + out = self.bn3(out) |
| 56 | + if self.downsample is not None: |
| 57 | + identity = self.downsample(x) |
| 58 | + out += identity |
| 59 | + return out |
| 60 | + |
| 61 | + |
| 62 | +class IResNet(nn.Module): |
| 63 | + fc_scale = 7 * 7 |
| 64 | + |
| 65 | + def __init__(self, block, layers, bn_norm, dropout=0, zero_init_residual=False, |
| 66 | + groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): |
| 67 | + super().__init__() |
| 68 | + self.inplanes = 64 |
| 69 | + self.dilation = 1 |
| 70 | + self.fp16 = fp16 |
| 71 | + if replace_stride_with_dilation is None: |
| 72 | + replace_stride_with_dilation = [False, False, False] |
| 73 | + if len(replace_stride_with_dilation) != 3: |
| 74 | + raise ValueError("replace_stride_with_dilation should be None " |
| 75 | + "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) |
| 76 | + self.groups = groups |
| 77 | + self.base_width = width_per_group |
| 78 | + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) |
| 79 | + self.bn1 = get_norm(bn_norm, self.inplanes) |
| 80 | + self.prelu = nn.PReLU(self.inplanes) |
| 81 | + self.layer1 = self._make_layer(block, 64, layers[0], bn_norm, stride=2) |
| 82 | + self.layer2 = self._make_layer(block, |
| 83 | + 128, |
| 84 | + layers[1], |
| 85 | + bn_norm, |
| 86 | + stride=2, |
| 87 | + dilate=replace_stride_with_dilation[0]) |
| 88 | + self.layer3 = self._make_layer(block, |
| 89 | + 256, |
| 90 | + layers[2], |
| 91 | + bn_norm, |
| 92 | + stride=2, |
| 93 | + dilate=replace_stride_with_dilation[1]) |
| 94 | + self.layer4 = self._make_layer(block, |
| 95 | + 512, |
| 96 | + layers[3], |
| 97 | + bn_norm, |
| 98 | + stride=2, |
| 99 | + dilate=replace_stride_with_dilation[2]) |
| 100 | + self.bn2 = get_norm(bn_norm, 512 * block.expansion) |
| 101 | + self.dropout = nn.Dropout(p=dropout, inplace=True) |
| 102 | + |
| 103 | + for m in self.modules(): |
| 104 | + if isinstance(m, nn.Conv2d): |
| 105 | + nn.init.normal_(m.weight, 0, 0.1) |
| 106 | + elif m.__class__.__name__.find('Norm') != -1: |
| 107 | + nn.init.constant_(m.weight, 1) |
| 108 | + nn.init.constant_(m.bias, 0) |
| 109 | + |
| 110 | + if zero_init_residual: |
| 111 | + for m in self.modules(): |
| 112 | + if isinstance(m, IBasicBlock): |
| 113 | + nn.init.constant_(m.bn2.weight, 0) |
| 114 | + |
| 115 | + def _make_layer(self, block, planes, blocks, bn_norm, stride=1, dilate=False): |
| 116 | + downsample = None |
| 117 | + previous_dilation = self.dilation |
| 118 | + if dilate: |
| 119 | + self.dilation *= stride |
| 120 | + stride = 1 |
| 121 | + if stride != 1 or self.inplanes != planes * block.expansion: |
| 122 | + downsample = nn.Sequential( |
| 123 | + conv1x1(self.inplanes, planes * block.expansion, stride), |
| 124 | + get_norm(bn_norm, planes * block.expansion), |
| 125 | + ) |
| 126 | + layers = [] |
| 127 | + layers.append( |
| 128 | + block(self.inplanes, planes, bn_norm, stride, downsample, self.groups, |
| 129 | + self.base_width, previous_dilation)) |
| 130 | + self.inplanes = planes * block.expansion |
| 131 | + for _ in range(1, blocks): |
| 132 | + layers.append( |
| 133 | + block(self.inplanes, |
| 134 | + planes, |
| 135 | + bn_norm, |
| 136 | + groups=self.groups, |
| 137 | + base_width=self.base_width, |
| 138 | + dilation=self.dilation)) |
| 139 | + |
| 140 | + return nn.Sequential(*layers) |
| 141 | + |
| 142 | + def forward(self, x): |
| 143 | + x = self.conv1(x) |
| 144 | + x = self.bn1(x) |
| 145 | + x = self.prelu(x) |
| 146 | + x = self.layer1(x) |
| 147 | + x = self.layer2(x) |
| 148 | + x = self.layer3(x) |
| 149 | + x = self.layer4(x) |
| 150 | + x = self.bn2(x) |
| 151 | + x = self.dropout(x) |
| 152 | + return x |
| 153 | + |
| 154 | + |
| 155 | +@BACKBONE_REGISTRY.register() |
| 156 | +def build_iresnet_backbone(cfg): |
| 157 | + """ |
| 158 | + Create a IResNet instance from config. |
| 159 | + Returns: |
| 160 | + ResNet: a :class:`ResNet` instance. |
| 161 | + """ |
| 162 | + |
| 163 | + # fmt: off |
| 164 | + bn_norm = cfg.MODEL.BACKBONE.NORM |
| 165 | + depth = cfg.MODEL.BACKBONE.DEPTH |
| 166 | + dropout = cfg.MODEL.BACKBONE.DROPOUT |
| 167 | + fp16 = cfg.SOLVER.AMP.ENABLED |
| 168 | + # fmt: on |
| 169 | + |
| 170 | + num_blocks_per_stage = { |
| 171 | + '18x': [2, 2, 2, 2], |
| 172 | + '34x': [3, 4, 6, 3], |
| 173 | + '50x': [3, 4, 14, 3], |
| 174 | + '100x': [3, 13, 30, 3], |
| 175 | + '200x': [6, 26, 60, 6], |
| 176 | + }[depth] |
| 177 | + |
| 178 | + model = IResNet(IBasicBlock, num_blocks_per_stage, bn_norm, dropout, fp16=fp16) |
| 179 | + return model |
0 commit comments