Compute loss for model. If both `labels` and `mask` are None, it degenerates to SimCLR unsupervised loss: https://arxiv.org/pdf/2002.05709.pdf Args: features: hidden vector of shape [bsz, n_views, ...]. labels: ground truth of shape [bsz].
(self, features, labels=None, mask=None)
| 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 |
| 93 | anchor_dot_contrast = torch.div( |
| 94 | torch.matmul(anchor_feature, contrast_feature.T), |
| 95 | self.temperature) |
| 96 | |
| 97 | # for numerical stability |
| 98 | logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) |
| 99 | logits = anchor_dot_contrast - logits_max.detach() |
| 100 | |
| 101 | # tile mask |
| 102 | mask = mask.repeat(anchor_count, contrast_count) |
| 103 | # mask-out self-contrast cases |