| 241 | + ', balance_weight=' + str(self.balance_weight) + ')' |
| 242 | |
| 243 | class Softmax(nn.Module): |
| 244 | def __init__(self, in_features, out_features): |
| 245 | super(Softmax, self).__init__() |
| 246 | self.classifier = nn.Linear(in_features, out_features, bias=False) |
| 247 | self.in_features = in_features |
| 248 | self.out_features = out_features |
| 249 | self.ce = torch.nn.CrossEntropyLoss() |
| 250 | self.reset_parameters() |
| 251 | |
| 252 | def reset_parameters(self): |
| 253 | nn.init.normal_(self.classifier.weight, std=0.001) |
| 254 | if self.classifier.bias: |
| 255 | nn.init.constant_(self.classifier.bias, 0.0) |
| 256 | |
| 257 | def forward(self, inputs, targets): |
| 258 | output = self.classifier(inputs) |
| 259 | loss = self.ce(output, targets) |
| 260 | top1 = accuracy(output.data, targets.data, topk=(1, 5))[0] |
| 261 | return {'loss': loss, 'logits': output, 'top1': top1} |
| 262 | |
| 263 | def __repr__(self): |
| 264 | return self.__class__.__name__ + '(' \ |
| 265 | + 'in_features=' + str(self.in_features) \ |
| 266 | + ', out_features=' + str(self.out_features) + ')' |
| 267 | |
| 268 | class Softmax_TripletLoss(nn.Module): |
| 269 | def __init__(self, in_features, out_features, tri_margin=0.3, balance_weight=1): |