# encoding: utf-8 """ @author: liaoxingyu @contact: sherlockliao01@gmail.com """ from torch import nn from .backbones.resnet import ResNet 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 Baseline(nn.Module): in_planes = 2048 def __init__(self, num_classes, last_stride, model_path): super(Baseline, self).__init__() self.base = ResNet(last_stride) self.base.load_param(model_path) self.gap = nn.AdaptiveAvgPool2d(1) # self.gap = nn.AdaptiveMaxPool2d(1) self.num_classes = num_classes self.bottleneck = nn.BatchNorm1d(self.in_planes) self.bottleneck.bias.requires_grad_(False) # no shift self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False) self.bottleneck.apply(weights_init_kaiming) self.classifier.apply(weights_init_classifier) def forward(self, x): global_feat = self.gap(self.base(x)) # (b, 2048, 1, 1) global_feat = global_feat.view(global_feat.shape[0], -1) # flatten to (bs, 2048) feat = self.bottleneck(global_feat) # normalize for angular softmax if self.training: cls_score = self.classifier(feat) return cls_score, global_feat # global feature for triplet loss else: return feat