def __init__(self, num_classes, camera_num, view_num, cfg, factory):
super(build_transformer, 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.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.SIE_CAMERA:
camera_num = camera_num
else:
camera_num = 0
if cfg.MODEL.SIE_VIEW:
view_num = view_num
else:
view_num = 0
self.base = factory[cfg.MODEL.TRANSFORMER_TYPE](
img_size=cfg.INPUT.SIZE_TRAIN,
sie_xishu=cfg.MODEL.SIE_COE,
camera=camera_num,
view=view_num,
stride_size=cfg.MODEL.STRIDE_SIZE,
drop_path_rate=cfg.MODEL.DROP_PATH,
drop_rate=cfg.MODEL.DROP_OUT,
attn_drop_rate=cfg.MODEL.ATT_DROP_RATE)
if cfg.MODEL.TRANSFORMER_TYPE == 'deit_small_patch16_224_TransReID':
self.in_planes = 384
if pretrain_choice == 'imagenet':
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(
model_path))
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 __init__(self, num_classes, camera_num, view_num, cfg, factory,
rearrange):
super(build_transformer_local, self).__init__()
model_path = cfg.MODEL.PRETRAIN_PATH
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.SIE_CAMERA:
camera_num = camera_num
else:
camera_num = 0
if cfg.MODEL.SIE_VIEW:
view_num = view_num
else:
view_num = 0
self.base = factory[cfg.MODEL.TRANSFORMER_TYPE](
img_size=cfg.INPUT.SIZE_TRAIN,
sie_xishu=cfg.MODEL.SIE_COE,
local_feature=cfg.MODEL.JPM,
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.shuffle_groups = cfg.MODEL.SHUFFLE_GROUP
print('using shuffle_groups size:{}'.format(self.shuffle_groups))
self.shift_num = cfg.MODEL.SHIFT_NUM
print('using shift_num size:{}'.format(self.shift_num))
self.divide_length = cfg.MODEL.DEVIDE_LENGTH
print('using divide_length size:{}'.format(self.divide_length))
self.rearrange = rearrange
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.cos_layer = cfg.MODEL.COS_LAYER
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 == '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)
print('using se_resnet101_ibn_a as a backbone')
elif model_name == 'resnet101_ibn_b':
self.in_planes = 2048
self.base = resnet101_ibn_b(last_stride)
print('using resnet101_ibn_b as a backbone')
elif model_name == 'efficientnet_b4':
print('using efficientnet_b2 as a backbone')
self.base = EfficientNet.from_pretrained('efficientnet-b2', advprop=False)
self.in_planes = self.base._fc.in_features
elif model_name == 'efficientnet_b7':
print('using efficientnet_b7 as a backbone')
self.base = EfficientNet.from_pretrained('efficientnet-b7', advprop=True)
self.in_planes = self.base._fc.in_features
elif model_name == 'resnext101_ibn_a':
self.in_planes = 2048
self.base = resnext101_ibn_a(last_stride)
print('using resnext101_ibn_a as a backbone')
elif model_name == 'HRnet':
self.in_planes = 2048
self.base = get_cls_net(cfg, pretrained = model_path)
else:
print('unsupported backbone! but got {}'.format(model_name))
if pretrain_choice == 'imagenet' and model_name != 'efficientnet_b4' and model_name != 'efficientnet_b7' and model_name != 'HRnet':
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(model_path))
if cfg.MODEL.POOLING_METHOD == 'GeM':
print('using GeM pooling')
self.gap = GeM()
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)
if self.model_name == 'HRnet':
self.attention_tconv = nn.Conv1d(self.in_planes, 1, 3, padding=1)
self.upsample0 = nn.Sequential(
nn.Conv2d(32, self.in_planes, kernel_size=1, stride=1, bias=False),
)
self.upsample1 = nn.Sequential(
nn.Conv2d(64, self.in_planes, kernel_size=1, stride=1, bias=False),
)
self.upsample2 = nn.Sequential(
nn.Conv2d(128, self.in_planes, kernel_size=1, stride=1, bias=False),
)
self.upsample3 = nn.Sequential(
nn.Conv2d(256, self.in_planes, kernel_size=1, stride=1, bias=False),
)
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.cos_layer = cfg.MODEL.COS_LAYER
self.neck = cfg.MODEL.NECK
self.neck_feat = cfg.TEST.NECK_FEAT
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 == '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)
print('using se_resnet101_ibn_a as a backbone')
elif model_name == 'resnet101_ibn_b':
self.in_planes = 2048
self.base = resnet101_ibn_b(last_stride)
print('using resnet101_ibn_b as a backbone')
elif 'efficientnet' in model_name:
self.base = EfficientNet.from_pretrained(model_name)
print('using {} as a backbone'.format(model_name))
self.in_planes = self.base._fc.in_features
else:
print('unsupported backbone! but got {}'.format(model_name))
#如果用efficentnet则这几句不要加
if pretrain_choice == 'imagenet':
#加载分布式权重时才需要加
# self.base = nn.DataParallel(self.base)
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(
model_path))
if cfg.MODEL.POOLING_METHOD == 'GeM':
print('using GeM pooling')
self.gap = GeM()
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 __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.cos_layer = cfg.MODEL.COS_LAYER
self.neck = cfg.MODEL.NECK
self.neck_feat = cfg.TEST.NECK_FEAT
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')
# resnet_ibn 与 常规的resnet区别在于在Bottleneck内使用了InstanceNorm.
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 == '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')
# se_resnet_ibn 在 se_resnet_ibn 基础上对Bottleneck增加了SE-layer对通道做attention操作
elif model_name == 'se_resnet101_ibn_a':
self.in_planes = 2048
self.base = se_resnet101_ibn_a(last_stride)
print('using se_resnet101_ibn_a as a backbone')
# resnet_ibn_b 与 resnet_ibn_a 的区别：
# resnet_ibn_a在Bottleneck内输出前将通道数拆分成两部分，一部分做InstanceNorm，一部分做batchNorm,然后拼接
# resnet_ibn_b将con1后面的batchNorm替换成InstanceNorm,并在在Bottleneck内输出前先做batchNorm再做InstanceNorm。
elif model_name == 'resnet101_ibn_b':
self.in_planes = 2048
self.base = resnet101_ibn_b(last_stride)
print('using resnet101_ibn_b as a backbone')
else:
print('unsupported backbone! but got {}'.format(model_name))
if pretrain_choice == 'imagenet':
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(
model_path))
if cfg.MODEL.POOLING_METHOD == 'GeM':
print('using GeM pooling')
self.gap = GeM()
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 __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')
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)
if pretrain_choice == 'imagenet' and model_name != 'efficientnet_b7':
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))
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))
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)
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 __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
self.model_name = 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.layernorm = cfg.MODEL.LAYERN0RM
if self.layernorm:
self.ln = nn.LayerNorm([cfg.MODEL.FEAT_SIZE])
if model_name == 'resnet50':
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.base = resnet50_ibn_a(last_stride)
print('using resnet50_ibn_a as a backbone')
elif model_name == 'resnet101_ibn_a':
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.base = se_resnet101_ibn_a(last_stride)
print('using se_resnet101_ibn_a as a backbone')
elif model_name == 'resnet101_ibn_b':
self.base = resnet101_ibn_b(last_stride)
print('using resnet101_ibn_b as a backbone')
elif model_name == 'efficientnet_b8':
self.base = EfficientNet.from_pretrained('efficientnet-b8')
elif model_name == 'efficientnet_b0':
self.base = EfficientNet.from_pretrained('efficientnet-b0')
elif model_name == 'efficientnet_b1':
self.base = EfficientNet.from_pretrained('efficientnet-b1')
elif model_name == 'efficientnet_b2':
self.base = EfficientNet.from_pretrained('efficientnet-b2')
elif model_name == 'efficientnet_b3':
self.base = EfficientNet.from_pretrained('efficientnet-b3')
elif model_name == 'efficientnet_b4':
self.base = EfficientNet.from_pretrained('efficientnet-b4')
elif model_name == 'efficientnet_b5':
self.base = EfficientNet.from_pretrained('efficientnet-b5')
elif model_name == 'efficientnet_b6':
self.base = EfficientNet.from_pretrained('efficientnet-b6')
elif model_name == 'efficientnet_b7':
self.base = EfficientNet.from_pretrained('efficientnet-b7')
elif model_name == 'densenet121': # feat size 1024
self.base = densenet121(pretrained=True)
elif model_name == 'densenet161':
self.base = densenet161(pretrained=True)
elif model_name == 'densenet169':
self.base = densenet169(pretrained=True)
elif model_name == 'nasnet':
self.base = ft_net_NAS() # feat size 4032
elif model_name == 'osnet':
self.base = osnet.build_osnet_backbone(cfg.MODEL.PRETRAIN_PATH)
elif model_name == 'sknet':
self.base = SKNet101()
elif model_name == 'cbam':
self.base = resnet101_cbam()
elif model_name == 'non_local':
self.base = Non_Local_101(last_stride)
elif model_name == 'inceptionv4': # feat size 1536
self.base = inception()
elif model_name == 'mgn':
self.base = MGN_res50()
else:
print('unsupported backbone! but got {}'.format(model_name))
self.in_planes = cfg.MODEL.FEAT_SIZE
if pretrain_choice == 'imagenet' and cfg.MODEL.PRETRAIN_PATH and 'resnet101' in cfg.MODEL.NAME:
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(
model_path))
if cfg.MODEL.POOLING_METHOD == 'GeM':
print('using GeM pooling')
self.gap = GeM()
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 == 'curricularface':
print('using {} loss with s:{}, m: {}'.format(
self.ID_LOSS_TYPE, cfg.SOLVER.COSINE_SCALE,
cfg.SOLVER.COSINE_MARGIN))
self.classifier = CurricularFace(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 __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
self.cfg = cfg
self.model_name = 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 = 1280
# model_weight_b0 = EfficientNet.from_pretrained('efficientnet-b0')
# model_weight_b0.to('cuda')
# mm = nn.Sequential(*model_weight_b0.named_children())
# self.base = model_weight_b0.extract_features
#
# from IPython import embed
# embed()
#print('using efficientnet-b0 as a backbone')
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 == '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)
print('using se_resnet101_ibn_a as a backbone')
elif model_name == 'se_resnet50_ibn_a':
self.in_planes = 2048
self.base = se_resnet50_ibn_a(last_stride)
print('using se_resnet101_ibn_a as a backbone')
elif model_name == 'resnet101_ibn_b':
self.in_planes = 2048
self.base = resnet101_ibn_b(last_stride)
print('using resnet101_ibn_b as a backbone')
elif model_name == 'resnet50_ibn_b':
self.in_planes = 2048
self.base = resnet50_ibn_b(last_stride)
print('using resnet50_ibn_b 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 == 'resnest50_ibn':
self.in_planes = 2048
self.base = resnest50_ibn(last_stride)
print('using resnest50_ibn 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 == 'resnest101_ibn':
self.in_planes = 2048
self.base = resnest101_ibn(last_stride)
print('using resnest101_ibn as a backbone')
elif model_name == 'efficientnet_b7':
# self.in_planes = 1280
#
# model_weight_b0 = EfficientNet.from_pretrained('efficientnet-b0')
# model_weight_b0.to('cuda')
# self.base = model_weight_b0.extract_features
self.base = EfficientNet.from_pretrained('efficientnet-b0')
self.in_planes = self.base._fc.in_features
print('using efficientnet-b0 as a backbone')
else:
print('unsupported backbone! but got {}'.format(model_name))
if pretrain_choice == 'imagenet' and model_name != 'efficientnet_b7':
# if model_name == 'efficientnet_b7':
# state_dict = torch.load(model_path)
# # self.base.load_state_dict(state_dict)
# if 'state_dict' in state_dict:
# param_dict = state_dict['state_dict']
# for i in param_dict:
# if 'fc' in i:
# continue
# self.state_dict()[i.replace('module.', '')].copy_(param_dict[i])
# else:
self.base.load_param(model_path)
print('Loading pretrained ImageNet model......from {}'.format(
model_path))
if cfg.MODEL.POOLING_METHOD == 'GeM':
print('using GeM pooling')
self.gap = GeM()
else:
self.gap = nn.AdaptiveAvgPool2d(1)
if cfg.MODEL.IF_USE_PCB:
self.pcb = PCB(cfg,
256,
num_classes,
0.5,
self.in_planes,
cut_at_pooling=False)
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)
