mirror of https://github.com/JDAI-CV/fast-reid.git
123 lines
4.7 KiB
Python
123 lines
4.7 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: xingyu liao
|
|
@contact: sherlockliao01@gmail.com
|
|
"""
|
|
|
|
from collections import namedtuple
|
|
|
|
from torch import nn
|
|
|
|
from fastreid.layers import get_norm, SELayer
|
|
from fastreid.modeling.backbones import BACKBONE_REGISTRY
|
|
|
|
|
|
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
|
"""3x3 convolution with padding"""
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
|
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
|
|
|
|
|
class bottleneck_IR(nn.Module):
|
|
def __init__(self, in_channel, depth, bn_norm, stride, with_se=False):
|
|
super(bottleneck_IR, self).__init__()
|
|
if in_channel == depth:
|
|
self.shortcut_layer = nn.MaxPool2d(1, stride)
|
|
else:
|
|
self.shortcut_layer = nn.Sequential(
|
|
nn.Conv2d(in_channel, depth, (1, 1), stride, bias=False),
|
|
get_norm(bn_norm, depth))
|
|
self.res_layer = nn.Sequential(
|
|
get_norm(bn_norm, in_channel),
|
|
nn.Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
|
|
nn.PReLU(depth),
|
|
nn.Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
|
|
get_norm(bn_norm, depth),
|
|
SELayer(depth, 16) if with_se else nn.Identity()
|
|
)
|
|
|
|
def forward(self, x):
|
|
shortcut = self.shortcut_layer(x)
|
|
res = self.res_layer(x)
|
|
return res + shortcut
|
|
|
|
|
|
class Bottleneck(namedtuple("Block", ["in_channel", "depth", "bn_norm", "stride", "with_se"])):
|
|
"""A named tuple describing a ResNet block."""
|
|
|
|
|
|
def get_block(in_channel, depth, bn_norm, num_units, with_se, stride=2):
|
|
return [Bottleneck(in_channel, depth, bn_norm, stride, with_se)] + \
|
|
[Bottleneck(depth, depth, bn_norm, 1, with_se) for _ in range(num_units - 1)]
|
|
|
|
|
|
def get_blocks(bn_norm, with_se, num_layers):
|
|
if num_layers == "50x":
|
|
blocks = [
|
|
get_block(in_channel=64, depth=64, bn_norm=bn_norm, num_units=3, with_se=with_se),
|
|
get_block(in_channel=64, depth=128, bn_norm=bn_norm, num_units=4, with_se=with_se),
|
|
get_block(in_channel=128, depth=256, bn_norm=bn_norm, num_units=14, with_se=with_se),
|
|
get_block(in_channel=256, depth=512, bn_norm=bn_norm, num_units=3, with_se=with_se)
|
|
]
|
|
elif num_layers == "100x":
|
|
blocks = [
|
|
get_block(in_channel=64, depth=64, bn_norm=bn_norm, num_units=3, with_se=with_se),
|
|
get_block(in_channel=64, depth=128, bn_norm=bn_norm, num_units=13, with_se=with_se),
|
|
get_block(in_channel=128, depth=256, bn_norm=bn_norm, num_units=30, with_se=with_se),
|
|
get_block(in_channel=256, depth=512, bn_norm=bn_norm, num_units=3, with_se=with_se)
|
|
]
|
|
elif num_layers == "152x":
|
|
blocks = [
|
|
get_block(in_channel=64, depth=64, bn_norm=bn_norm, num_units=3, with_se=with_se),
|
|
get_block(in_channel=64, depth=128, bn_norm=bn_norm, num_units=8, with_se=with_se),
|
|
get_block(in_channel=128, depth=256, bn_norm=bn_norm, num_units=36, with_se=with_se),
|
|
get_block(in_channel=256, depth=512, bn_norm=bn_norm, num_units=3, with_se=with_se)
|
|
]
|
|
return blocks
|
|
|
|
|
|
class ResNetIR(nn.Module):
|
|
def __init__(self, num_layers, bn_norm, drop_ratio, with_se):
|
|
super(ResNetIR, self).__init__()
|
|
assert num_layers in ["50x", "100x", "152x"], "num_layers should be 50,100, or 152"
|
|
blocks = get_blocks(bn_norm, with_se, num_layers)
|
|
self.input_layer = nn.Sequential(nn.Conv2d(3, 64, (3, 3), 1, 1, bias=False),
|
|
get_norm(bn_norm, 64),
|
|
nn.PReLU(64))
|
|
self.output_layer = nn.Sequential(get_norm(bn_norm, 512),
|
|
nn.Dropout(drop_ratio))
|
|
modules = []
|
|
for block in blocks:
|
|
for bottleneck in block:
|
|
modules.append(
|
|
bottleneck_IR(bottleneck.in_channel,
|
|
bottleneck.depth,
|
|
bottleneck.bn_norm,
|
|
bottleneck.stride,
|
|
bottleneck.with_se))
|
|
self.body = nn.Sequential(*modules)
|
|
|
|
def forward(self, x):
|
|
x = self.input_layer(x)
|
|
x = self.body(x)
|
|
x = self.output_layer(x)
|
|
return x
|
|
|
|
|
|
@BACKBONE_REGISTRY.register()
|
|
def build_resnetIR_backbone(cfg):
|
|
"""
|
|
Create a ResNetIR instance from config.
|
|
Returns:
|
|
ResNet: a :class:`ResNet` instance.
|
|
"""
|
|
|
|
# fmt: off
|
|
bn_norm = cfg.MODEL.BACKBONE.NORM
|
|
with_se = cfg.MODEL.BACKBONE.WITH_SE
|
|
depth = cfg.MODEL.BACKBONE.DEPTH
|
|
# fmt: on
|
|
|
|
model = ResNetIR(depth, bn_norm, 0.5, with_se)
|
|
return model
|