| from pytorch_metric_learning.distances import CosineSimilarity |
| import torch |
|
|
|
|
| class InfoNCELoss(): |
| def __init__(self, device, k, temperature=0.07, threshold=1.0, fna=False): |
| super(InfoNCELoss, self).__init__() |
| self.device = device |
| self.similarity = CosineSimilarity() |
| self.k = k |
| self.temperature = temperature |
| self.threshold = threshold |
| self.fna = fna |
|
|
| def __call__(self, x, y, labels, sts): |
| false_negatives = self.detect_false_negative(sts) |
| indices_tuple = self.get_all_pairs_indices(labels, false_negatives) |
|
|
| mat = self.similarity(x, y) |
| a1, p, a2, n = indices_tuple |
| pos_pair, neg_pair = [], [] |
| if len(a1) > 0: |
| pos_pair = mat[a1, p] |
| if len(a2) > 0: |
| neg_pair = mat[a2, n] |
|
|
| if len(neg_pair) > 0 and self.k > -1: |
| paired = list(zip(neg_pair.tolist(), a2.tolist(), n.tolist())) |
| selected = sorted(paired, key=lambda x: x[0], reverse=True)[:self.k] |
| _, x, y = list(zip(*selected)) |
| x = torch.tensor(x).to(a2.device) |
| y = torch.tensor(y).to(n.device) |
|
|
| neg_pair = mat[x, y] |
| indices_tuple = (a1, p, x, y) |
|
|
| return self._compute_loss(pos_pair, neg_pair, indices_tuple), len(pos_pair) |
|
|
| def detect_false_negative(self, embs): |
| mat = torch.matmul(embs, torch.t(embs)) |
| return torch.where(mat >= self.threshold) |
|
|
| def _compute_loss(self, pos_pairs, neg_pairs, indices_tuple): |
| a1, p, a2, _ = indices_tuple |
|
|
| if len(a1) > 0 and len(a2) > 0: |
| dtype = neg_pairs.dtype |
|
|
| if not self.similarity.is_inverted: |
| pos_pairs = -pos_pairs |
| neg_pairs = -neg_pairs |
|
|
| pos_pairs = pos_pairs.unsqueeze(1) / self.temperature |
| neg_pairs = neg_pairs / self.temperature |
| n_per_p = a2.unsqueeze(0) == a1.unsqueeze(1) |
| neg_pairs = neg_pairs * n_per_p |
| neg_pairs[n_per_p == 0] = torch.finfo(dtype).min |
|
|
| max_val = torch.max( |
| pos_pairs, torch.max(neg_pairs, dim=1, keepdim=True)[0] |
| ).detach() |
| numerator = torch.exp(pos_pairs - max_val).squeeze(1) |
| denominator = torch.sum(torch.exp(neg_pairs - max_val), dim=1) + numerator |
| log_exp = torch.log((numerator / denominator) + torch.finfo(dtype).tiny) |
| return torch.mean(-log_exp) |
|
|
| return 0 |
|
|
| def get_all_pairs_indices(self, labels, false_negatives): |
| labels1 = labels.unsqueeze(1) |
| labels2 = labels.unsqueeze(0) |
| matches = (labels1 == labels2).byte() |
| diffs = matches ^ 1 |
|
|
| diffs[false_negatives[0], false_negatives[1]] = 0 |
| if self.fna: |
| matches[false_negatives[0], false_negatives[1]] = 1 |
|
|
| diffs.fill_diagonal_(0) |
| matches.fill_diagonal_(1) |
|
|
| a1_idx, p_idx = torch.where(matches) |
| a2_idx, n_idx = torch.where(diffs) |
| return a1_idx, p_idx, a2_idx, n_idx |
|
|
