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)
neg_weight = 1 / num_neg if num_neg > 0 else 0
if weight is not None:
match_shape(weight, num_pos)
pos_loss = F.binary_cross_entropy_with_logits(
pos_scores,
pos_scores.new_ones(()).expand(num_pos),
reduction="sum",
weight=weight,
)
neg_loss = F.binary_cross_entropy_with_logits(
neg_scores,
neg_scores.new_zeros(()).expand(num_pos, num_neg),
reduction="sum",
weight=weight.unsqueeze(-1) if weight is not None else None,
)
loss = pos_loss + neg_weight * neg_loss
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)
neg_weight = 1 / num_neg if num_neg > 0 else 0
pos_loss = F.binary_cross_entropy_with_logits(pos_scores,
pos_scores.new_ones(
()).expand(num_pos),
reduction="sum")
neg_loss = F.binary_cross_entropy_with_logits(
neg_scores,
neg_scores.new_zeros(()).expand(num_pos, num_neg),
reduction="sum",
)
loss = pos_loss + neg_weight * neg_loss
return loss