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