Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. It also supports the unsupervised contrastive loss in SimCLR From: https://github.com/HobbitLong/SupContrast
| 33 | from project_utils.k_means_utils import test_kmeans_semi_sup |
| 34 | |
| 35 | class SupConLoss(torch.nn.Module): |
| 36 | """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. |
| 37 | It also supports the unsupervised contrastive loss in SimCLR |
| 38 | From: https://github.com/HobbitLong/SupContrast""" |
| 39 | def __init__(self, temperature=0.07, contrast_mode='all', |
| 40 | base_temperature=0.07): |
| 41 | super(SupConLoss, self).__init__() |
| 42 | self.temperature = temperature |
| 43 | self.contrast_mode = contrast_mode |
| 44 | self.base_temperature = base_temperature |
| 45 | |
| 46 | def forward(self, features, labels=None, mask=None): |
| 47 | """Compute loss for model. If both `labels` and `mask` are None, |
| 48 | it degenerates to SimCLR unsupervised loss: |
| 49 | https://arxiv.org/pdf/2002.05709.pdf |
| 50 | Args: |
| 51 | features: hidden vector of shape [bsz, n_views, ...]. |
| 52 | labels: ground truth of shape [bsz]. |
| 53 | mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j |
| 54 | has the same class as sample i. Can be asymmetric. |
| 55 | Returns: |
| 56 | A loss scalar. |
| 57 | """ |
| 58 | device = (torch.device('cuda') |
| 59 | if features.is_cuda |
| 60 | else torch.device('cpu')) |
| 61 | |
| 62 | if len(features.shape) < 3: |
| 63 | raise ValueError('`features` needs to be [bsz, n_views, ...],' |
| 64 | 'at least 3 dimensions are required') |
| 65 | if len(features.shape) > 3: |
| 66 | features = features.view(features.shape[0], features.shape[1], -1) |
| 67 | |
| 68 | batch_size = features.shape[0] |
| 69 | if labels is not None and mask is not None: |
| 70 | raise ValueError('Cannot define both `labels` and `mask`') |
| 71 | elif labels is None and mask is None: |
| 72 | mask = torch.eye(batch_size, dtype=torch.float32).to(device) |
| 73 | elif labels is not None: |
| 74 | labels = labels.contiguous().view(-1, 1) |
| 75 | if labels.shape[0] != batch_size: |
| 76 | raise ValueError('Num of labels does not match num of features') |
| 77 | mask = torch.eq(labels, labels.T).float().to(device) |
| 78 | else: |
| 79 | mask = mask.float().to(device) |
| 80 | |
| 81 | contrast_count = features.shape[1] |
| 82 | contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0) |
| 83 | if self.contrast_mode == 'one': |
| 84 | anchor_feature = features[:, 0] |
| 85 | anchor_count = 1 |
| 86 | elif self.contrast_mode == 'all': |
| 87 | anchor_feature = contrast_feature |
| 88 | anchor_count = contrast_count |
| 89 | else: |
| 90 | raise ValueError('Unknown mode: {}'.format(self.contrast_mode)) |
| 91 | |
| 92 | # compute logits |