| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | class SoftmaxLoss(nn.Module): |
| 10 | def __init__(self, cfg): |
| 11 | super(SoftmaxLoss, self).__init__() |
| 12 | |
| 13 | self.feat_dim = cfg.MODEL.EMBEDDING_DIM |
| 14 | self.num_classes = cfg.MODEL.LOSS.LUT_SIZE |
| 15 | |
| 16 | self.bottleneck = nn.BatchNorm1d(self.feat_dim) |
| 17 | self.bottleneck.bias.requires_grad_(False) # no shift |
| 18 | self.classifier = nn.Linear(self.feat_dim, self.num_classes, bias=False) |
| 19 | |
| 20 | self.bottleneck.apply(weights_init_kaiming) |
| 21 | self.classifier.apply(weights_init_classifier) |
| 22 | |
| 23 | def forward(self, inputs, labels): |
| 24 | """ |
| 25 | Args: |
| 26 | inputs: feature matrix with shape (batch_size, feat_dim). |
| 27 | labels: ground truth labels with shape (num_classes). |
| 28 | """ |
| 29 | assert inputs.size(0) == labels.size(0), "features.size(0) is not equal to labels.size(0)" |
| 30 | |
| 31 | target = labels.clone() |
| 32 | target[target >= self.num_classes] = 5554 |
| 33 | |
| 34 | feat = self.bottleneck(inputs) |
| 35 | score = self.classifier(feat) |
| 36 | loss = F.cross_entropy(score, target, ignore_index=5554) |
| 37 | |
| 38 | return loss |
| 39 | |
| 40 | |
| 41 | def weights_init_kaiming(m): |