from typing import Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
[docs]
def categorical_crossentropy_loss(y_true: torch.Tensor, y_pred: torch.Tensor, from_logits: bool = True) -> torch.Tensor:
"""
Compute categorical crossentropy
:param y_true: torch.Tensor, ground truth
:param y_pred: torch.Tensor, model output
:param from_logits: bool, `True` means y_pred has not transformed by softmax, default True
:return: torch.Tensor, loss value
"""
if from_logits:
return -(F.log_softmax(y_pred, dim=1) * y_true).sum(dim=1)
return -(torch.log(y_pred, dim=1) * y_true).sum(dim=1)
[docs]
def cosine_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 20.0) -> torch.Tensor:
"""
Compute cosine loss
:param y_true: torch.Tensor, ground truth.
The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair.
:param y_pred: torch.Tensor, model output.
The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair.
:param tau: float, scale factor, default 20
:return: torch.Tensor, loss value
""" # NOQA
# modified from: https://github.com/bojone/CoSENT/blob/124c368efc8a4b179469be99cb6e62e1f2949d39/cosent.py#L79
y_true = y_true[::2, 0]
y_true = (y_true[:, None] < y_true[None, :]).float()
y_pred = F.normalize(y_pred, p=2, dim=1)
y_pred = torch.sum(y_pred[::2] * y_pred[1::2], dim=1) * tau
y_pred = y_pred[:, None] - y_pred[None, :]
y_pred = (y_pred - (1 - y_true) * 1e12).view(-1)
zero = torch.Tensor([0]).to(y_pred.device)
y_pred = torch.concat((zero, y_pred), dim=0)
return torch.logsumexp(y_pred, dim=0)
[docs]
def angle_loss(y_true: torch.Tensor, y_pred: torch.Tensor, tau: float = 1.0, pooling_strategy: str = 'sum'):
"""
Compute angle loss
:param y_true: torch.Tensor, ground truth.
The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair.
:param y_pred: torch.Tensor, model output.
The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair.
:param tau: float, scale factor, default 1.0
:return: torch.Tensor, loss value
""" # NOQA
y_true = y_true[::2, 0]
y_true = (y_true[:, None] < y_true[None, :]).float()
y_pred_re, y_pred_im = torch.chunk(y_pred, 2, dim=1)
a = y_pred_re[::2]
b = y_pred_im[::2]
c = y_pred_re[1::2]
d = y_pred_im[1::2]
# (a+bi) / (c+di)
# = ((a+bi) * (c-di)) / ((c+di) * (c-di))
# = ((ac + bd) + i(bc - ad)) / (c^2 + d^2)
# = (ac + bd) / (c^2 + d^2) + i(bc - ad)/(c^2 + d^2)
z = torch.sum(c**2 + d**2, dim=1, keepdim=True)
re = (a * c + b * d) / z
im = (b * c - a * d) / z
dz = torch.sum(a**2 + b**2, dim=1, keepdim=True)**0.5
dw = torch.sum(c**2 + d**2, dim=1, keepdim=True)**0.5
re /= (dz / dw)
im /= (dz / dw)
y_pred = torch.concat((re, im), dim=1)
if pooling_strategy == 'sum':
pooling = torch.sum(y_pred, dim=1)
elif pooling_strategy == 'mean':
pooling = torch.mean(y_pred, dim=1)
else:
raise ValueError(f'Unsupported pooling strategy: {pooling_strategy}')
y_pred = torch.abs(pooling) * tau # absolute delta angle
y_pred = y_pred[:, None] - y_pred[None, :]
y_pred = (y_pred - (1 - y_true) * 1e12).view(-1)
zero = torch.Tensor([0]).to(y_pred.device)
y_pred = torch.concat((zero, y_pred), dim=0)
return torch.logsumexp(y_pred, dim=0)
[docs]
def in_batch_negative_loss(
y_true: torch.Tensor,
y_pred: torch.Tensor,
tau: float = 20.0,
negative_weights: float = 0.0
) -> torch.Tensor:
"""
Compute in-batch negative loss, i.e., contrastive loss
:param y_true: torch.Tensor, ground truth.
The y_true must be zigzag style, such as [x[0][0], x[0][1], x[1][0], x[1][1], ...], where (x[0][0], x[0][1]) stands for a pair.
:param y_pred: torch.Tensor, model output.
The y_pred must be zigzag style, such as [o[0][0], o[0][1], o[1][0], o[1][1], ...], where (o[0][0], o[0][1]) stands for a pair.
:param tau: float, scale factor, default 20.0
:param negative_weights: float, negative weights, default 0.0
:return: torch.Tensor, loss value
""" # NOQA
device = y_true.device
def make_target_matrix(y_true: torch.Tensor):
idxs = torch.arange(0, y_pred.shape[0]).int().to(device)
y_true = y_true.int()
idxs_1 = idxs[None, :]
idxs_2 = (idxs + 1 - idxs % 2 * 2)[:, None]
idxs_1 *= y_true.T
idxs_1 += (y_true.T == 0).int() * -2
idxs_2 *= y_true
idxs_2 += (y_true == 0).int() * -1
y_true = (idxs_1 == idxs_2).float()
return y_true
neg_mask = make_target_matrix(y_true == 0)
y_true = make_target_matrix(y_true)
# compute similarity
y_pred = F.normalize(y_pred, dim=1, p=2)
similarities = y_pred @ y_pred.T # dot product
similarities = similarities - torch.eye(y_pred.shape[0]).to(device) * 1e12
similarities = similarities * tau
if negative_weights > 0:
similarities += neg_mask * negative_weights
return categorical_crossentropy_loss(y_true, similarities, from_logits=True).mean()
[docs]
def contrastive_with_negative_loss(
text: torch.Tensor,
pos: torch.Tensor,
neg: Optional[torch.Tensor] = None,
tau: float = 20.0
) -> torch.Tensor:
"""
Compute contrastive with negative loss
:param text: torch.Tensor, text.
:param pos: torch.Tensor, positive samples of text.
:param neg: torch.Tensor, negative samples of text.
:param tau: float, scale factor, default 20.0
:return: torch.Tensor, loss value
"""
target = torch.cat((pos, neg), dim=0) if neg is not None else pos # (2B, D)
q_norm = torch.nn.functional.normalize(text, p=2, dim=1) # (B, D)
t_norm = torch.nn.functional.normalize(target, p=2, dim=1) # (2B, D)
scores = torch.mm(q_norm, t_norm.transpose(0, 1)) * tau # (B, 2B)
labels = torch.tensor(
range(len(scores)), dtype=torch.long, device=scores.device
)
return nn.CrossEntropyLoss()(scores, labels)