Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. It also supports the unsupervised contrastive loss in SimCLR
| 9 | |
| 10 | |
| 11 | class SupConLoss(nn.Module): |
| 12 | """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. |
| 13 | It also supports the unsupervised contrastive loss in SimCLR""" |
| 14 | def __init__(self, temperature=0.07, contrast_mode='all', |
| 15 | base_temperature=0.07): |
| 16 | super(SupConLoss, self).__init__() |
| 17 | self.temperature = temperature |
| 18 | self.contrast_mode = contrast_mode |
| 19 | self.base_temperature = base_temperature |
| 20 | |
| 21 | def forward(self, features, labels=None, mask=None): |
| 22 | """Compute loss for model. If both `labels` and `mask` are None, |
| 23 | it degenerates to SimCLR unsupervised loss: |
| 24 | https://arxiv.org/pdf/2002.05709.pdf |
| 25 | |
| 26 | Args: |
| 27 | features: hidden vector of shape [bsz, n_views, ...]. |
| 28 | labels: ground truth of shape [bsz]. |
| 29 | mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j |
| 30 | has the same class as sample i. Can be asymmetric. |
| 31 | Returns: |
| 32 | A loss scalar. |
| 33 | """ |
| 34 | device = (torch.device('cuda') |
| 35 | if features.is_cuda |
| 36 | else torch.device('cpu')) |
| 37 | |
| 38 | if len(features.shape) < 3: |
| 39 | raise ValueError('`features` needs to be [bsz, n_views, ...],' |
| 40 | 'at least 3 dimensions are required') |
| 41 | if len(features.shape) > 3: |
| 42 | features = features.view(features.shape[0], features.shape[1], -1) |
| 43 | |
| 44 | batch_size = features.shape[0] |
| 45 | if labels is not None and mask is not None: |
| 46 | raise ValueError('Cannot define both `labels` and `mask`') |
| 47 | elif labels is None and mask is None: |
| 48 | mask = torch.eye(batch_size, dtype=torch.float32).to(device) |
| 49 | elif labels is not None: |
| 50 | labels = labels.contiguous().view(-1, 1) |
| 51 | if labels.shape[0] != batch_size: |
| 52 | raise ValueError('Num of labels does not match num of features') |
| 53 | mask = torch.eq(labels, labels.T).float().to(device) |
| 54 | else: |
| 55 | mask = mask.float().to(device) |
| 56 | |
| 57 | contrast_count = features.shape[1] |
| 58 | contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0) |
| 59 | if self.contrast_mode == 'one': |
| 60 | anchor_feature = features[:, 0] |
| 61 | anchor_count = 1 |
| 62 | elif self.contrast_mode == 'all': |
| 63 | anchor_feature = contrast_feature |
| 64 | anchor_count = contrast_count |
| 65 | else: |
| 66 | raise ValueError('Unknown mode: {}'.format(self.contrast_mode)) |
| 67 | |
| 68 | # compute logits |