|
|
| import torch |
| import torch.nn as nn |
| import timm |
|
|
| class BRAINHybridLoss(nn.Module): |
| def __init__(self, margin=1.0): |
| super(BRAINHybridLoss, self).__init__() |
| self.margin = margin |
| self.ce_loss = nn.CrossEntropyLoss() |
|
|
| def forward(self, logits, labels, embeddings): |
| loss_ce = self.ce_loss(logits, labels) |
| |
| dist_mat = torch.cdist(embeddings, embeddings, p=2.0) |
| |
| labels_reshaped = labels.unsqueeze(1) |
| mask_pos = (labels_reshaped == labels_reshaped.T).float() |
| mask_neg = (labels_reshaped != labels_reshaped.T).float() |
| |
| negative_dist = dist_mat + (1.0 - mask_neg) * 1e6 |
| hard_negatives = negative_dist.min(dim=1)[0] |
| |
| positive_dist = dist_mat * mask_pos |
| hard_positives = positive_dist.max(dim=1)[0] |
| |
| loss_triplet = F.relu(hard_positives - hard_negatives + self.margin).mean() |
| |
| return loss_ce + (0.1 * loss_triplet) |
|
|
