def __init__(self, cfg):
super().__init__()
self._cfg = cfg
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
self.register_buffer(
"pixel_mean",
torch.Tensor(cfg.MODEL.PIXEL_MEAN).view(1, -1, 1, 1))
self.register_buffer(
"pixel_std",
torch.Tensor(cfg.MODEL.PIXEL_STD).view(1, -1, 1, 1))
# fmt: off
# backbone
bn_norm = cfg.MODEL.BACKBONE.NORM
num_splits = cfg.MODEL.BACKBONE.NORM_SPLIT
with_se = cfg.MODEL.BACKBONE.WITH_SE
# fmt :on
backbone = build_backbone(cfg)
self.backbone = nn.Sequential(backbone.conv1, backbone.bn1,
backbone.relu, backbone.maxpool,
backbone.layer1, backbone.layer2,
backbone.layer3[0])
res_conv4 = nn.Sequential(*backbone.layer3[1:])
res_g_conv5 = backbone.layer4
res_p_conv5 = nn.Sequential(
Bottleneck(1024,
512,
bn_norm,
num_splits,
False,
with_se,
downsample=nn.Sequential(
nn.Conv2d(1024, 2048, 1, bias=False),
get_norm(bn_norm, 2048, num_splits))),
Bottleneck(2048, 512, bn_norm, num_splits, False, with_se),
Bottleneck(2048, 512, bn_norm, num_splits, False, with_se))
res_p_conv5.load_state_dict(backbone.layer4.state_dict())
# branch1
self.b1 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_g_conv5))
self.b1_head = build_reid_heads(cfg)
# branch2
self.b2 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
self.b2_head = build_reid_heads(cfg)
self.b21_head = build_reid_heads(cfg)
self.b22_head = build_reid_heads(cfg)
# branch3
self.b3 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
self.b3_head = build_reid_heads(cfg)
self.b31_head = build_reid_heads(cfg)
self.b32_head = build_reid_heads(cfg)
self.b33_head = build_reid_heads(cfg)
def __init__(self,
backbone,
num_classes,
last_stride,
with_ibn,
gcb,
stage_with_gcb,
pretrain=True,
model_path=''):
super().__init__()
self.num_classes = num_classes
if 'resnet' in backbone:
self.base = ResNet.from_name(backbone, last_stride, with_ibn, gcb,
stage_with_gcb)
self.base.load_pretrain(model_path)
self.in_planes = 2048
elif 'osnet' in backbone:
if with_ibn:
self.base = osnet_ibn_x1_0(pretrained=pretrain)
else:
self.base = osnet_x1_0(pretrained=pretrain)
self.in_planes = 512
else:
print(f'not support {backbone} backbone')
# global branch
self.global_reduction = nn.Sequential(
nn.Conv2d(self.in_planes, 512, 1), nn.BatchNorm2d(512),
nn.ReLU(True))
self.gap = nn.AdaptiveAvgPool2d(1)
self.global_bn = bn2d_no_bias(512)
self.global_classifier = nn.Linear(512, self.num_classes, bias=False)
# mask brach
self.part = Bottleneck(2048, 512)
self.batch_drop = BatchDrop(1.0, 0.33)
self.part_pool = nn.AdaptiveMaxPool2d(1)
self.part_reduction = nn.Sequential(nn.Conv2d(self.in_planes, 1024, 1),
nn.BatchNorm2d(1024),
nn.ReLU(True))
self.part_bn = bn2d_no_bias(1024)
self.part_classifier = nn.Linear(1024, self.num_classes, bias=False)
# initialize
self.part.apply(weights_init_kaiming)
self.global_reduction.apply(weights_init_kaiming)
self.part_reduction.apply(weights_init_kaiming)
self.global_classifier.apply(weights_init_classifier)
self.part_classifier.apply(weights_init_classifier)
def __init__(self, cfg):
super().__init__()
self.register_buffer(
"pixel_mean",
torch.Tensor(cfg.MODEL.PIXEL_MEAN).view(1, -1, 1, 1))
self.register_buffer(
"pixel_std",
torch.Tensor(cfg.MODEL.PIXEL_STD).view(1, -1, 1, 1))
self._cfg = cfg
# backbone
bn_norm = cfg.MODEL.BACKBONE.NORM
num_splits = cfg.MODEL.BACKBONE.NORM_SPLIT
with_se = cfg.MODEL.BACKBONE.WITH_SE
backbone = build_backbone(cfg)
self.backbone = nn.Sequential(backbone.conv1, backbone.bn1,
backbone.relu, backbone.maxpool,
backbone.layer1, backbone.layer2,
backbone.layer3[0])
res_conv4 = nn.Sequential(*backbone.layer3[1:])
res_g_conv5 = backbone.layer4
res_p_conv5 = nn.Sequential(
Bottleneck(1024,
512,
bn_norm,
num_splits,
False,
with_se,
downsample=nn.Sequential(
nn.Conv2d(1024, 2048, 1, bias=False),
get_norm(bn_norm, 2048, num_splits))),
Bottleneck(2048, 512, bn_norm, num_splits, False, with_se),
Bottleneck(2048, 512, bn_norm, num_splits, False, with_se))
res_p_conv5.load_state_dict(backbone.layer4.state_dict())
if cfg.MODEL.HEADS.POOL_LAYER == 'avgpool':
pool_layer = nn.AdaptiveAvgPool2d(1)
elif cfg.MODEL.HEADS.POOL_LAYER == 'maxpool':
pool_layer = nn.AdaptiveMaxPool2d(1)
elif cfg.MODEL.HEADS.POOL_LAYER == 'gempool':
pool_layer = GeneralizedMeanPoolingP()
else:
pool_layer = nn.Identity()
# head
in_feat = cfg.MODEL.HEADS.IN_FEAT
num_classes = cfg.MODEL.HEADS.NUM_CLASSES
# branch1
self.b1 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_g_conv5))
self.b1_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b1_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
# branch2
self.b2 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
self.b2_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b2_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
self.b21_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b21_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
self.b22_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b22_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
# branch3
self.b3 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
self.b3_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b3_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
self.b31_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b31_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
self.b32_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b32_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
self.b33_pool = self._build_pool_reduce(pool_layer, reduce_dim=in_feat)
self.b33_head = build_reid_heads(cfg, in_feat, num_classes,
nn.Identity())
class BDNet(nn.Module):
def __init__(self,
backbone,
num_classes,
last_stride,
with_ibn,
gcb,
stage_with_gcb,
pretrain=True,
model_path=''):
super().__init__()
self.num_classes = num_classes
if 'resnet' in backbone:
self.base = ResNet.from_name(backbone, last_stride, with_ibn, gcb,
stage_with_gcb)
self.base.load_pretrain(model_path)
self.in_planes = 2048
elif 'osnet' in backbone:
if with_ibn:
self.base = osnet_ibn_x1_0(pretrained=pretrain)
else:
self.base = osnet_x1_0(pretrained=pretrain)
self.in_planes = 512
else:
print(f'not support {backbone} backbone')
# global branch
self.global_reduction = nn.Sequential(
nn.Conv2d(self.in_planes, 512, 1), nn.BatchNorm2d(512),
nn.ReLU(True))
self.gap = nn.AdaptiveAvgPool2d(1)
self.global_bn = bn2d_no_bias(512)
self.global_classifier = nn.Linear(512, self.num_classes, bias=False)
# mask brach
self.part = Bottleneck(2048, 512)
self.batch_drop = BatchDrop(1.0, 0.33)
self.part_pool = nn.AdaptiveMaxPool2d(1)
self.part_reduction = nn.Sequential(nn.Conv2d(self.in_planes, 1024, 1),
nn.BatchNorm2d(1024),
nn.ReLU(True))
self.part_bn = bn2d_no_bias(1024)
self.part_classifier = nn.Linear(1024, self.num_classes, bias=False)
# initialize
self.part.apply(weights_init_kaiming)
self.global_reduction.apply(weights_init_kaiming)
self.part_reduction.apply(weights_init_kaiming)
self.global_classifier.apply(weights_init_classifier)
self.part_classifier.apply(weights_init_classifier)
def forward(self, x, label=None):
# feature extractor
feat = self.base(x)
# global branch
g_feat = self.global_reduction(feat)
g_feat = self.gap(g_feat) # (bs, 512, 1, 1)
g_bn_feat = self.global_bn(g_feat) # (bs, 512, 1, 1)
g_bn_feat = g_bn_feat.view(-1, g_bn_feat.shape[1]) # (bs, 512)
# mask branch
p_feat = self.part(feat)
p_feat = self.batch_drop(p_feat)
p_feat = self.part_pool(p_feat) # (bs, 512, 1, 1)
p_feat = self.part_reduction(p_feat)
p_bn_feat = self.part_bn(p_feat)
p_bn_feat = p_bn_feat.view(-1, p_bn_feat.shape[1]) # (bs, 512)
if self.training:
global_cls = self.global_classifier(g_bn_feat)
part_cls = self.part_classifier(p_bn_feat)
return global_cls, part_cls, g_feat.view(
-1, g_feat.shape[1]), p_feat.view(-1, p_feat.shape[1])
return torch.cat([g_bn_feat, p_bn_feat], dim=1)
def load_params_wo_fc(self, state_dict):
state_dict.pop('global_classifier.weight')
state_dict.pop('part_classifier.weight')
res = self.load_state_dict(state_dict, strict=False)
print(f'missing keys {res.missing_keys}')
# assert str(res.missing_keys) == str(['classifier.weight',]), 'issue loading pretrained weights'
def unfreeze_all_layers(self, ):
self.train()
for p in self.parameters():
p.requires_grad = True
def unfreeze_specific_layer(self, names):
if isinstance(names, str):
names = [names]
for name, module in self.named_children():
if name in names:
module.train()
for p in module.parameters():
p.requires_grad = True
else:
module.eval()
for p in module.parameters():
p.requires_grad = False
def from_config(cls, cfg):
bn_norm = cfg.MODEL.BACKBONE.NORM
with_se = cfg.MODEL.BACKBONE.WITH_SE
all_blocks = build_backbone(cfg)
# backbone
backbone = nn.Sequential(all_blocks.conv1, all_blocks.bn1,
all_blocks.relu, all_blocks.maxpool,
all_blocks.layer1, all_blocks.layer2,
all_blocks.layer3[0])
res_conv4 = nn.Sequential(*all_blocks.layer3[1:])
res_g_conv5 = all_blocks.layer4
res_p_conv5 = nn.Sequential(
Bottleneck(1024,
512,
bn_norm,
False,
with_se,
downsample=nn.Sequential(
nn.Conv2d(1024, 2048, 1, bias=False),
get_norm(bn_norm, 2048))),
Bottleneck(2048, 512, bn_norm, False, with_se),
Bottleneck(2048, 512, bn_norm, False, with_se))
res_p_conv5.load_state_dict(all_blocks.layer4.state_dict())
# branch
neck1 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_g_conv5))
b1_head = build_heads(cfg)
# branch2
neck2 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
b2_head = build_heads(cfg)
b21_head = build_heads(cfg)
b22_head = build_heads(cfg)
# branch3
neck3 = nn.Sequential(copy.deepcopy(res_conv4),
copy.deepcopy(res_p_conv5))
b3_head = build_heads(cfg)
b31_head = build_heads(cfg)
b32_head = build_heads(cfg)
b33_head = build_heads(cfg)
return {
'backbone': backbone,
'neck1': neck1,
'neck2': neck2,
'neck3': neck3,
'b1_head': b1_head,
'b2_head': b2_head,
'b21_head': b21_head,
'b22_head': b22_head,
'b3_head': b3_head,
'b31_head': b31_head,
'b32_head': b32_head,
'b33_head': b33_head,
'pixel_mean': cfg.MODEL.PIXEL_MEAN,
'pixel_std': cfg.MODEL.PIXEL_STD,
'loss_kwargs': {
# loss name
'loss_names': cfg.MODEL.LOSSES.NAME,
# loss hyperparameters
'ce': {
'eps': cfg.MODEL.LOSSES.CE.EPSILON,
'alpha': cfg.MODEL.LOSSES.CE.ALPHA,
'scale': cfg.MODEL.LOSSES.CE.SCALE
},
'tri': {
'margin': cfg.MODEL.LOSSES.TRI.MARGIN,
'norm_feat': cfg.MODEL.LOSSES.TRI.NORM_FEAT,
'hard_mining': cfg.MODEL.LOSSES.TRI.HARD_MINING,
'scale': cfg.MODEL.LOSSES.TRI.SCALE
},
'circle': {
'margin': cfg.MODEL.LOSSES.CIRCLE.MARGIN,
'gamma': cfg.MODEL.LOSSES.CIRCLE.GAMMA,
'scale': cfg.MODEL.LOSSES.CIRCLE.SCALE
},
'cosface': {
'margin': cfg.MODEL.LOSSES.COSFACE.MARGIN,
'gamma': cfg.MODEL.LOSSES.COSFACE.GAMMA,
'scale': cfg.MODEL.LOSSES.COSFACE.SCALE
}
}
}