import torch import torch.nn as nn from .backbones.resnet_ibn_a import resnet50_ibn_a,resnet101_ibn_a from .backbones.resnext_ibn import resnext101_ibn_a from .layers.pooling import GeM, GeneralizedMeanPooling,GeneralizedMeanPoolingP import torch.nn.functional as F # from efficientnet_pytorch import EfficientNet import copy def weights_init_kaiming(m): classname = m.__class__.__name__ if classname.find('Linear') != -1: nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out') nn.init.constant_(m.bias, 0.0) elif classname.find('Conv') != -1: nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in') if m.bias is not None: nn.init.constant_(m.bias, 0.0) elif classname.find('BatchNorm') != -1: if m.affine: nn.init.constant_(m.weight, 1.0) nn.init.constant_(m.bias, 0.0) def weights_init_classifier(m): classname = m.__class__.__name__ if classname.find('Linear') != -1: nn.init.normal_(m.weight, std=0.001) if m.bias: nn.init.constant_(m.bias, 0.0) class Backbone(nn.Module): def __init__(self, num_classes, cfg): super(Backbone, self).__init__() last_stride = cfg.MODEL.LAST_STRIDE model_path = cfg.MODEL.PRETRAIN_PATH model_name = cfg.MODEL.NAME pretrain_choice = cfg.MODEL.PRETRAIN_CHOICE self.neck = cfg.MODEL.NECK self.neck_feat = cfg.TEST.NECK_FEAT self.model_name = model_name if model_name == 'resnet50': self.in_planes = 2048 self.base = ResNet(last_stride=last_stride, block=Bottleneck, frozen_stages=cfg.MODEL.FROZEN, layers=[3, 4, 6, 3]) print('using resnet50 as a backbone') elif model_name == 'resnet50_ibn_a': self.in_planes = 2048 self.base = resnet50_ibn_a(last_stride) print('using resnet50_ibn_a as a backbone') elif model_name == 'resnet152': self.in_planes = 2048 self.base = ResNet(last_stride=last_stride, block=Bottleneck, layers=[3, 8, 36, 3]) elif model_name == 'resnet101_ibn_a': self.in_planes = 2048 self.base = resnet101_ibn_a(last_stride, frozen_stages=cfg.MODEL.FROZEN) print('using resnet101_ibn_a as a backbone') elif model_name == 'se_resnet101_ibn_a': self.in_planes = 2048 self.base = se_resnet101_ibn_a(last_stride,frozen_stages=cfg.MODEL.FROZEN) print('using se_resnet101_ibn_a as a backbone') elif model_name == 'efficientnet_b7': print('using efficientnet_b7 as a backbone') self.base = EfficientNet.from_pretrained('efficientnet-b7', advprop=False) self.in_planes = self.base._fc.in_features elif model_name == 'densenet169_ibn_a': self.in_planes = 1664 self.base = densenet169_ibn_a() print('using densenet169_ibn_a as a backbone') elif model_name == 'resnest50': self.in_planes = 2048 self.base = resnest50(last_stride) print('using resnest50 as a backbone') elif model_name == 'resnest101': self.in_planes = 2048 self.base = resnest101(last_stride) print('using resnest101 as a backbone') elif model_name == 'resnest200': self.in_planes = 2048 self.base = resnest200(last_stride) print('using resnest200 as a backbone') elif model_name == 'resnest269': self.in_planes = 2048 self.base = resnest269(last_stride) print('using resnest269 as a backbone') elif model_name == 'resnext101_ibn_a': self.in_planes = 2048 self.base = resnext101_ibn_a() print('using resnext101_ibn_a as a backbone') elif model_name == 'hrnet32': self.in_planes = 2048 self.base = hrnet_32(last_stride) print('using hrnet32 as a backbone') else: print('unsupported backbone! but got {}'.format(model_name)) if cfg.MODEL.POOLING_METHOD == 'gempoolP': print('using GeMP pooling') self.gap = GeneralizedMeanPoolingP() elif cfg.MODEL.POOLING_METHOD == 'gempool': print('using GeM pooling') self.gap = GeM(freeze_p=False) else: self.gap = nn.AdaptiveAvgPool2d(1) self.num_classes = num_classes self.ID_LOSS_TYPE = cfg.MODEL.ID_LOSS_TYPE if self.ID_LOSS_TYPE == 'arcface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Arcface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'cosface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Cosface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'amsoftmax': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = AMSoftmax(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'circle': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE, cfg.SOLVER.COSINE_SCALE, cfg.SOLVER.COSINE_MARGIN)) self.classifier = CircleLoss(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) else: self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier.apply(weights_init_classifier) self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) self.bottleneck.apply(weights_init_kaiming) def forward(self, x, label=None,cam_label=None): # label is unused if self.cos_layer == 'no' if self.model_name =='efficientnet_b7': x = self.base.extract_features(x) else: x = self.base(x) global_feat = nn.functional.avg_pool2d(x, x.shape[2:4]) global_feat = global_feat.view(global_feat.shape[0], -1) # flatten to (bs, 2048) if self.neck == 'no': feat = global_feat elif self.neck == 'bnneck': feat = self.bottleneck(global_feat) # print("$$$$$$",feat) if self.training: if self.ID_LOSS_TYPE in ('arcface', 'cosface', 'amsoftmax', 'circle'): cls_score = self.classifier(feat, label) else: cls_score = self.classifier(feat) return cls_score, global_feat # global feature for triplet loss else: if self.neck_feat == 'after': # print("Test with feature after BN") return feat else: # print("Test with feature before BN") return global_feat def load_param(self, trained_path): param_dict = torch.load(trained_path,map_location='cpu') if 'state_dict' in param_dict: param_dict = param_dict['state_dict'] for i in param_dict: if 'classifier' in i or 'arcface' in i: continue self.state_dict()[i.replace('module.','')].copy_(param_dict[i]) print('Loading pretrained model from {}'.format(trained_path)) def load_un_param(self, trained_path): param_dict = torch.load(trained_path,map_location='cpu') if 'state_dict' in param_dict: param_dict = param_dict['state_dict'] for i in self.state_dict(): if 'classifier' in i or 'arcface' in i: continue self.state_dict()[i].copy_(param_dict[i]) print('Loading pretrained model from {}'.format(trained_path)) def load_param_finetune(self, model_path): param_dict = torch.load(model_path) for i in param_dict: self.state_dict()[i].copy_(param_dict[i]) print('Loading pretrained model for finetuning from {}'.format(model_path)) class build_transformer(nn.Module): def __init__(self, num_classes, camera_num, view_num, cfg, factory): super(build_transformer, self).__init__() model_path = cfg.MODEL.PRETRAIN_PATH model_name = cfg.MODEL.NAME pretrain_choice = cfg.MODEL.PRETRAIN_CHOICE self.cos_layer = cfg.MODEL.COS_LAYER self.neck = cfg.MODEL.NECK self.neck_feat = cfg.TEST.NECK_FEAT print('using Transformer_type: {} as a backbone'.format(cfg.MODEL.Transformer_TYPE)) if cfg.MODEL.CAMERA_EMBEDDING: camera_num = camera_num else: camera_num = 0 if cfg.MODEL.VIEWPOINT_EMBEDDING: view_num = view_num else: view_num = 0 self.base = factory[cfg.MODEL.Transformer_TYPE](img_size=cfg.INPUT.SIZE_TRAIN, aie_xishu=cfg.MODEL.AIE_COE,local_feature=cfg.MODEL.LOCAL_F, camera=camera_num, view=view_num, stride_size=cfg.MODEL.STRIDE_SIZE, drop_path_rate=cfg.MODEL.DROP_PATH) if pretrain_choice == 'imagenet': self.base.load_param(model_path) print('Loading pretrained ImageNet model......from {}'.format(model_path)) elif pretrain_choice == 'self': param_dict = torch.load(model_path, map_location='cpu') for i in param_dict: if 'classifier' in i: continue self.state_dict()[i].copy_(param_dict[i]) print('Loading finetune model......from {}'.format(model_path)) self.gap = nn.AdaptiveAvgPool2d(1) self.num_classes = num_classes self.in_planes = self.base.embed_dim self.ID_LOSS_TYPE = cfg.MODEL.ID_LOSS_TYPE if self.ID_LOSS_TYPE == 'arcface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Arcface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'cosface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Cosface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'amsoftmax': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = AMSoftmax(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'circle': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE, cfg.SOLVER.COSINE_SCALE, cfg.SOLVER.COSINE_MARGIN)) self.classifier = CircleLoss(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) else: self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier.apply(weights_init_classifier) self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) self.bottleneck.apply(weights_init_kaiming) def forward(self, x, label=None, cam_label= None): # label is unused if self.cos_layer == 'no' global_feat = self.base(x, cam_label=cam_label) feat = self.bottleneck(global_feat) if self.training: if self.ID_LOSS_TYPE in ('arcface', 'cosface', 'amsoftmax', 'circle'): cls_score = self.classifier(feat, label) else: cls_score = self.classifier(feat) return cls_score, global_feat # global feature for triplet loss else: if self.neck_feat == 'after': # print("Test with feature after BN") return feat else: # print("Test with feature before BN") return global_feat def load_param(self, trained_path): param_dict = torch.load(trained_path,map_location='cpu') for i in param_dict: if 'classifier' in i or 'arcface' in i or 'gap' in i: continue # self.state_dict()[i].copy_(param_dict[i]) self.state_dict()[i.replace('module.', '')].copy_(param_dict[i]) print('Loading pretrained model from {}'.format(trained_path)) class build_transformer_local(nn.Module): def __init__(self, num_classes, camera_num, view_num, cfg, factory): super(build_transformer_local, self).__init__() model_path = cfg.MODEL.PRETRAIN_PATH model_name = cfg.MODEL.NAME pretrain_choice = cfg.MODEL.PRETRAIN_CHOICE self.cos_layer = cfg.MODEL.COS_LAYER self.neck = cfg.MODEL.NECK self.neck_feat = cfg.TEST.NECK_FEAT self.in_planes = 768 print('using Transformer_type: {} as a backbone'.format(cfg.MODEL.Transformer_TYPE)) if cfg.MODEL.CAMERA_EMBEDDING: camera_num = camera_num else: camera_num = 0 if cfg.MODEL.VIEWPOINT_EMBEDDING: view_num = view_num else: view_num = 0 self.base = factory[cfg.MODEL.Transformer_TYPE](img_size=cfg.INPUT.SIZE_TRAIN, aie_xishu=cfg.MODEL.AIE_COE, local_feature=cfg.MODEL.LOCAL_F, camera=camera_num, view=view_num, stride_size=cfg.MODEL.STRIDE_SIZE, drop_path_rate=cfg.MODEL.DROP_PATH) if pretrain_choice == 'imagenet': self.base.load_param(model_path) print('Loading pretrained ImageNet model......from {}'.format(model_path)) block = self.base.blocks[-1] layer_norm = self.base.norm self.b1 = nn.Sequential( copy.deepcopy(block), copy.deepcopy(layer_norm) ) self.b2 = nn.Sequential( copy.deepcopy(block), copy.deepcopy(layer_norm) ) self.num_classes = num_classes self.ID_LOSS_TYPE = cfg.MODEL.ID_LOSS_TYPE if self.ID_LOSS_TYPE == 'arcface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Arcface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'cosface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Cosface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'amsoftmax': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = AMSoftmax(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'circle': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE, cfg.SOLVER.COSINE_SCALE, cfg.SOLVER.COSINE_MARGIN)) self.classifier = CircleLoss(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) else: self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier.apply(weights_init_classifier) self.classifier_1 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_1.apply(weights_init_classifier) self.classifier_2 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_2.apply(weights_init_classifier) self.classifier_3 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_3.apply(weights_init_classifier) self.classifier_4 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_4.apply(weights_init_classifier) self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) self.bottleneck.apply(weights_init_kaiming) self.bottleneck_1 = nn.BatchNorm1d(self.in_planes) self.bottleneck_1.bias.requires_grad_(False) self.bottleneck_1.apply(weights_init_kaiming) self.bottleneck_2 = nn.BatchNorm1d(self.in_planes) self.bottleneck_2.bias.requires_grad_(False) self.bottleneck_2.apply(weights_init_kaiming) self.bottleneck_3 = nn.BatchNorm1d(self.in_planes) self.bottleneck_3.bias.requires_grad_(False) self.bottleneck_3.apply(weights_init_kaiming) self.bottleneck_4 = nn.BatchNorm1d(self.in_planes) self.bottleneck_4.bias.requires_grad_(False) self.bottleneck_4.apply(weights_init_kaiming) self.groups = cfg.MODEL.PATCH_SHUFFLE print('using group size:{}'.format(self.groups)) self.yiwei = cfg.MODEL.YIWEI_NUM print('using yiwei size:{}'.format(self.yiwei)) self.divide_length = cfg.MODEL.DEVIDE_LENGTH print('using divide_length size:{}'.format(self.divide_length)) if pretrain_choice == 'self': param_dict = torch.load(model_path, map_location='cpu') for i in param_dict: if 'classifier' in i: continue self.state_dict()[i].copy_(param_dict[i]) print('Loading finetune model......from {}'.format(model_path)) def forward(self, x, label=None, cam_label= None): # label is unused if self.cos_layer == 'no' features = self.base(x, cam_label=cam_label) batchsize = features.size(0) dim = features.size(-1) # branch1 b1_feat = self.b1(features) # [64, 129, 768] global_feat = b1_feat[:, 0] # branch2 token = features[:, 0] feature_length = features.size(1) - 1 patch_length = feature_length // self.divide_length feature_random = torch.cat([features[:, self.yiwei:], features[:, 1:self.yiwei]], dim=1) channels_per_group = feature_length // self.groups x = feature_random x = x.view(batchsize, self.groups, channels_per_group, dim) x = torch.transpose(x, 1, 2).contiguous() x = x.view(batchsize, -1, dim) b1_local_feat = x[:, :patch_length] b1_local_feat = self.b2(torch.cat((token.unsqueeze(1), b1_local_feat), dim=1)) local_feat_1 = b1_local_feat[:, 0] # branch3 b2_local_feat = x[:, patch_length:patch_length*2] b2_local_feat = self.b2(torch.cat((token.unsqueeze(1), b2_local_feat), dim=1)) local_feat_2 = b2_local_feat[:, 0] # branch4 b3_local_feat = x[:, patch_length*2:patch_length*3] b3_local_feat = self.b2(torch.cat((token.unsqueeze(1), b3_local_feat), dim=1)) local_feat_3 = b3_local_feat[:, 0] # branch5 b4_local_feat = x[:, patch_length*3:patch_length*4] b4_local_feat = self.b2(torch.cat((token.unsqueeze(1), b4_local_feat), dim=1)) local_feat_4 = b4_local_feat[:, 0] feat = self.bottleneck(global_feat) local_feat_1_bn = self.bottleneck_1(local_feat_1) local_feat_2_bn = self.bottleneck_2(local_feat_2) local_feat_3_bn = self.bottleneck_3(local_feat_3) local_feat_4_bn = self.bottleneck_4(local_feat_4) if self.training: if self.ID_LOSS_TYPE in ('arcface', 'cosface', 'amsoftmax', 'circle'): cls_score = self.classifier(feat, label) else: cls_score = self.classifier(feat) cls_score_1 = self.classifier_1(local_feat_1_bn) cls_score_2 = self.classifier_2(local_feat_2_bn) cls_score_3 = self.classifier_3(local_feat_3_bn) cls_score_4 = self.classifier_4(local_feat_4_bn) return [cls_score, cls_score_1, cls_score_2, cls_score_3, cls_score_4 ], [global_feat, local_feat_1, local_feat_2, local_feat_3, local_feat_4] # global feature for triplet loss else: if self.neck_feat == 'after': # print('hhhh') return torch.cat( [feat, local_feat_1_bn / 4, local_feat_2_bn / 4, local_feat_3_bn / 4, local_feat_4_bn / 4], dim=1) # print('global') # return feat else: # print("Test with feature before BN") return global_feat def load_param(self, trained_path): param_dict = torch.load(trained_path, map_location = 'cpu') for i in param_dict: if 'classifier' in i or 'arcface' in i or 'bottleneck' in i or 'gap' in i: continue #self.state_dict()[i].copy_(param_dict[i]) self.state_dict()[i.replace('module.', '')].copy_(param_dict[i]) print('Loading pretrained model from {}'.format(trained_path)) class build_transformer_local_no_shuffle(nn.Module): def __init__(self, num_classes, camera_num, view_num, cfg, factory): super(build_transformer_local_no_shuffle, self).__init__() model_path = cfg.MODEL.PRETRAIN_PATH model_name = cfg.MODEL.NAME pretrain_choice = cfg.MODEL.PRETRAIN_CHOICE self.cos_layer = cfg.MODEL.COS_LAYER self.neck = cfg.MODEL.NECK self.neck_feat = cfg.TEST.NECK_FEAT self.in_planes = 768 print('using Transformer_type: {} as a backbone'.format(cfg.MODEL.Transformer_TYPE)) if cfg.MODEL.CAMERA_EMBEDDING: camera_num = camera_num else: camera_num = 0 if cfg.MODEL.VIEWPOINT_EMBEDDING: view_num = view_num else: view_num = 0 self.base = factory[cfg.MODEL.Transformer_TYPE](img_size=cfg.INPUT.SIZE_TRAIN, local_feature=cfg.MODEL.LOCAL_F, camera=camera_num, view=view_num, stride_size=cfg.MODEL.STRIDE_SIZE, drop_path_rate=cfg.MODEL.DROP_PATH) if pretrain_choice == 'imagenet': self.base.load_param(model_path, map_location='cpu') print('Loading pretrained ImageNet model......from {}'.format(model_path)) elif pretrain_choice == 'self': param_dict = torch.load(model_path, map_location='cpu') for i in param_dict: if 'classifier' in i: continue self.state_dict()[i].copy_(param_dict[i]) print('Loading finetune model......from {}'.format(model_path)) block = self.base.blocks[-1] layer_norm = self.base.norm self.b1 = nn.Sequential( copy.deepcopy(block), copy.deepcopy(layer_norm) ) self.b2 = nn.Sequential( copy.deepcopy(block), copy.deepcopy(layer_norm) ) self.num_classes = num_classes self.ID_LOSS_TYPE = cfg.MODEL.ID_LOSS_TYPE if self.ID_LOSS_TYPE == 'arcface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Arcface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'cosface': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = Cosface(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'amsoftmax': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE,cfg.SOLVER.COSINE_SCALE,cfg.SOLVER.COSINE_MARGIN)) self.classifier = AMSoftmax(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) elif self.ID_LOSS_TYPE == 'circle': print('using {} with s:{}, m: {}'.format(self.ID_LOSS_TYPE, cfg.SOLVER.COSINE_SCALE, cfg.SOLVER.COSINE_MARGIN)) self.classifier = CircleLoss(self.in_planes, self.num_classes, s=cfg.SOLVER.COSINE_SCALE, m=cfg.SOLVER.COSINE_MARGIN) else: self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier.apply(weights_init_classifier) self.classifier_1 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_1.apply(weights_init_classifier) self.classifier_2 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_2.apply(weights_init_classifier) self.classifier_3 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_3.apply(weights_init_classifier) self.classifier_4 = nn.Linear(self.in_planes, self.num_classes, bias=False) self.classifier_4.apply(weights_init_classifier) self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) self.bottleneck.apply(weights_init_kaiming) self.bottleneck_1 = nn.BatchNorm1d(self.in_planes) self.bottleneck_1.bias.requires_grad_(False) self.bottleneck_1.apply(weights_init_kaiming) self.bottleneck_2 = nn.BatchNorm1d(self.in_planes) self.bottleneck_2.bias.requires_grad_(False) self.bottleneck_2.apply(weights_init_kaiming) self.bottleneck_3 = nn.BatchNorm1d(self.in_planes) self.bottleneck_3.bias.requires_grad_(False) self.bottleneck_3.apply(weights_init_kaiming) self.bottleneck_4 = nn.BatchNorm1d(self.in_planes) self.bottleneck_4.bias.requires_grad_(False) self.bottleneck_4.apply(weights_init_kaiming) self.groups = cfg.MODEL.PATCH_SHUFFLE print('using group size:{}'.format(self.groups)) print('not using yiwei') def forward(self, x, label=None, cam_label= None): # label is unused if self.cos_layer == 'no' features = self.base(x, cam_label=cam_label) batchsize = features.size(0) dim = features.size(-1) # branch1 b1_feat = self.b1(features) # [64, 129, 768] global_feat = b1_feat[:, 0] # branch2 token = features[:, 0] feature_length = features.size(1) - 1 patch_length = feature_length // 4 x = features[:, 1:] b1_local_feat = x[:, :patch_length] b1_local_feat = self.b2(torch.cat((token.unsqueeze(1), b1_local_feat), dim=1)) local_feat_1 = b1_local_feat[:, 0] # branch3 b2_local_feat = x[:, patch_length:patch_length*2] b2_local_feat = self.b2(torch.cat((token.unsqueeze(1), b2_local_feat), dim=1)) local_feat_2 = b2_local_feat[:, 0] # branch4 b3_local_feat = x[:, patch_length*2:patch_length*3] b3_local_feat = self.b2(torch.cat((token.unsqueeze(1), b3_local_feat), dim=1)) local_feat_3 = b3_local_feat[:, 0] # branch5 b4_local_feat = x[:, patch_length*3:] b4_local_feat = self.b2(torch.cat((token.unsqueeze(1), b4_local_feat), dim=1)) local_feat_4 = b4_local_feat[:, 0] feat = self.bottleneck(global_feat) local_feat_1_bn = self.bottleneck_1(local_feat_1) local_feat_2_bn = self.bottleneck_2(local_feat_2) local_feat_3_bn = self.bottleneck_3(local_feat_3) local_feat_4_bn = self.bottleneck_4(local_feat_4) if self.training: if self.ID_LOSS_TYPE in ('arcface', 'cosface', 'amsoftmax', 'circle'): cls_score = self.classifier(feat, label) else: cls_score = self.classifier(feat) cls_score_1 = self.classifier_1(local_feat_1_bn) cls_score_2 = self.classifier_2(local_feat_2_bn) cls_score_3 = self.classifier_3(local_feat_3_bn) cls_score_4 = self.classifier_4(local_feat_4_bn) return [cls_score, cls_score_1, cls_score_2, cls_score_3, cls_score_4 ], [global_feat, local_feat_1, local_feat_2, local_feat_3, local_feat_4] # global feature for triplet loss else: if self.neck_feat == 'after': return torch.cat( [feat, local_feat_1_bn / 4, local_feat_2_bn / 4, local_feat_3_bn / 4, local_feat_4_bn / 4], dim=1) # return feat else: # print("Test with feature before BN") return global_feat def load_param(self, trained_path): param_dict = torch.load(trained_path, map_location='cpu') for i in param_dict: if 'classifier' in i or 'arcface' in i or 'bottleneck' in i or 'gap' in i: continue self.state_dict()[i].copy_(param_dict[i]) print('Loading pretrained model from {}'.format(trained_path)) def make_model(cfg, num_class, camera_num=0, view_num=0): if cfg.MODEL.NAME == 'transformer': if cfg.MODEL.LOCAL_F: if cfg.MODEL.NO_SHUFFLE: model = build_transformer_local_no_shuffle(num_class, camera_num, view_num, cfg, __factory_hh) print('===========building build_transformer_local_no_shuffle===========') else: model = build_transformer_local(num_class, camera_num, view_num, cfg, __factory_hh) print('===========building build_transformer_local===========') # model = build_transformer_local_8(num_class, camera_num, view_num, cfg, __factory_hh) # print('===========building build_transformer_local_8===========') else: model = build_transformer(num_class, camera_num, view_num, cfg, __factory_hh) print('===========building transformer===========') else: print('===========ResNet===========') model = Backbone(num_class, cfg) return model