def forward(
self,
pos_scores: FloatTensorType,
neg_scores: FloatTensorType,
weight: Optional[FloatTensorType],
) -> FloatTensorType:
num_pos = match_shape(pos_scores, -1)
num_neg = match_shape(neg_scores, num_pos, -1)
# FIXME Workaround for https://github.com/pytorch/pytorch/issues/15223.
if num_pos == 0 or num_neg == 0:
return torch.zeros((),
device=pos_scores.device,
requires_grad=True)
if weight is not None:
match_shape(weight, num_pos)
loss_per_sample = F.margin_ranking_loss(
neg_scores,
pos_scores.unsqueeze(1),
target=pos_scores.new_full((1, 1), -1, dtype=torch.float),
margin=self.margin,
reduction="none",
)
loss = (loss_per_sample * weight.unsqueeze(-1)).sum()
else:
# more memory efficient way if no weights
loss = F.margin_ranking_loss(
neg_scores,
pos_scores.unsqueeze(1),
target=pos_scores.new_full((1, 1), -1, dtype=torch.float),
margin=self.margin,
reduction="sum",
)
return loss
def forward(self, pos_scores: FloatTensorType,
neg_scores: FloatTensorType) -> FloatTensorType:
num_pos = match_shape(pos_scores, -1)
num_neg = match_shape(neg_scores, num_pos, -1)
# FIXME Workaround for https://github.com/pytorch/pytorch/issues/15223.
if num_pos == 0 or num_neg == 0:
return torch.zeros((),
device=pos_scores.device,
requires_grad=True)
loss = F.margin_ranking_loss(
neg_scores,
pos_scores.unsqueeze(1),
target=pos_scores.new_full((1, 1), -1, dtype=torch.float),
margin=self.margin,
reduction="sum",
)
return loss