Spaces:
Paused
Paused
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class MarginRankLoss(nn.Module): | |
| def __init__(self): | |
| super(MarginRankLoss, self).__init__() | |
| self.loss = 0.0 | |
| def forward(self, img_list, margin=0): | |
| length = len(img_list) | |
| self.loss = 0.0 | |
| B, C, H, W = img_list[0].shape | |
| for i in range(length - 1): | |
| for j in range(i + 1, length): | |
| self.loss = self.loss + torch.sum(F.relu(img_list[j].sum(-1).sum(-1).sum(-1) - img_list[i].sum(-1).sum(-1).sum(-1) + margin)) | |
| self.loss = self.loss / (B*length*(length-1)/2) | |
| return self.loss | |
| class RankLoss(nn.Module): | |
| def __init__(self): | |
| super(RankLoss, self).__init__() | |
| self.countloss_criterion = nn.MSELoss(reduction='sum') | |
| self.rankloss_criterion = MarginRankLoss() | |
| self.rankloss = 0.0 | |
| self.losses = {} | |
| def forward(self, unlabeled_results): | |
| self.rankloss = 0.0 | |
| self.losses = {} | |
| if unlabeled_results is None: | |
| self.rankloss = 0.0 | |
| elif isinstance(unlabeled_results, tuple) and len(unlabeled_results) > 0: | |
| self.rankloss = self.rankloss_criterion(unlabeled_results) | |
| elif isinstance(unlabeled_results, list) and len(unlabeled_results) > 0: | |
| count = 0 | |
| for i in range(len(unlabeled_results)): | |
| if isinstance(unlabeled_results[i], tuple) and len(unlabeled_results[i]) > 0: | |
| temploss = self.rankloss_criterion(unlabeled_results[i]) | |
| self.losses.update({'unlabel_{}_loss'.format(str(i+1)): temploss}) | |
| self.rankloss += temploss | |
| count += 1 | |
| if count > 0: | |
| self.rankloss = self.rankloss / count | |
| return self.rankloss | |