BCELoss with weights for labels (like weighted_cross_entropy_with_logits in TF)

[velikodniy]

I've implemented an analog of weighted_cross_entropy_with_logits in my current project. It's useful for working with imbalanced datasets. I want to add it to PyTorch but I'm in doubt if it is really needed for others.

For example, my implementation:

def weighted_binary_cross_entropy_with_logits(logits, targets, pos_weight, weight=None, size_average=True, reduce=True):
    if not (targets.size() == logits.size()):
        raise ValueError("Target size ({}) must be the same as input size ({})".format(targets.size(), logits.size()))

    max_val = (-logits).clamp(min=0)
    log_weight = 1 + (pos_weight - 1) * targets
    loss = (1 - targets) * logits + log_weight * (((-max_val).exp() + (-logits - max_val).exp()).log() + max_val)

    if weight is not None:
        loss = loss * weight

    if not reduce:
        return loss
    elif size_average:
        return loss.mean()
    else:
        return loss.sum()

class WeightedBCEWithLogitsLoss(torch.nn.Module):
    def __init__(self, pos_weight, weight=None, size_average=True, reduce=True):
        super().__init__()
        self.register_buffer('weight', weight)
        self.register_buffer('pos_weight', pos_weight)
        self.size_average = size_average
        self.reduce = reduce

    def forward(self, input, target):
        pos_weight = Variable(self.pos_weight) if not isinstance(self.pos_weight, Variable) else self.pos_weight
        if self.weight is not None:
            weight = Variable(self.weight) if not isinstance(self.weight, Variable) else self.weight
            return weighted_binary_cross_entropy_with_logits(input, target,
                                                             pos_weight,
                                                             weight=weight
                                                             size_average=self.size_average,
                                                             reduce=self.reduce)
        else:
            return weighted_binary_cross_entropy_with_logits(input, target,
                                                             pos_weight,
                                                             weight=None,
                                                             size_average=self.size_average,
                                                             reduce=self.reduce)

(Of course, tests and WeightedBCELoss should be written too.)

[velikodniy]

Another way is to add pos_weight argument to BCELoss and BCEWithLogits loss.

[soumith]

we discussed this proposal yesterday, and I dont think it makes sense to add it. target weights are simply a multiplier that one can do before-hand with *, and it's not that common of a use-case.

[velikodniy]

Note that pos_weight is multiplied only by the first addend in the formula for BCE loss. It's not the weight for the whole target. I cannot see the simple way to do it beforehand.

Proposed loss looks like:

pos_weight * targets * -log(sigmoid(logits)) + (1 - targets) * -log(1 - sigmoid(logits))

BCEWithLogitsLoss with multiplier m for targets:

m * targets * -log(sigmoid(logits)) + (1 - m * targets) * -log(1 - sigmoid(logits))

If these formulae are the same, m should be equal:

(log(1 - sigmoid(logits)) - pos_weight * log(sigmoid(logits)))/(log(1 - sigmoid(logits)) - log(sigmoid(logits)))

This formula is quite complex, contains logits, and I guess it is numerically unstable.

Maybe I didn't understand your idea.

[sbonner0]

I too would like this functionality in pytorch. I'm not exactly sure how one would do it using BCEWithLogitsLoss alone? I'm not sure it would be equivalent with tf.weighted_cross_entropy_with_logits

[soumith]

@velikodniy I should've re-visited this much earlier, but it slipped my queue. The introduction of pos_weight indeed makes it a bit more complicated.

I think we should add that option to BCELoss and BCEWithLogitsLoss directly.

[velikodniy]

@soumith Thank you for your reply. I agree with you about adding the option pos_weight to BCELoss/BCEWithLogitsLoss and their F. analogs. I will prepare a pool request soon.