torch/legacy/nn/MultiLabelSoftMarginCriterion.py

import torch
from .Criterion import Criterion
from .Sigmoid import Sigmoid
from .BCECriterion import BCECriterion


class MultiLabelSoftMarginCriterion(Criterion):
    """
    A MultiLabel multiclass criterion based on sigmoid:

    the loss is:
    l(x, y) = - sum_i y[i] * log(p[i]) + (1 - y[i]) * log (1 - p[i])
    where p[i] = exp(x[i]) / (1 + exp(x[i]))

    and with weights:
    l(x, y) = - sum_i weights[i] (y[i] * log(p[i]) + (1 - y[i]) * log (1 - p[i]))


    """

    def __init__(self, weights=None):
        super(MultiLabelSoftMarginCriterion, self).__init__()
        self.lsm = Sigmoid()
        self.nll = BCECriterion(weights)

    def updateOutput(self, input, target):
        input = input if input.nelement() == 1 else input.squeeze()
        target = target if target.nelement() == 1 else target.squeeze()
        self.lsm.updateOutput(input)
        self.nll.updateOutput(self.lsm.output, target)
        self.output = self.nll.output
        return self.output

    def updateGradInput(self, input, target):
        size = input.size()
        input = input if input.nelement() == 1 else input.squeeze()
        target = target if target.nelement() == 1 else target.squeeze()
        self.nll.updateGradInput(self.lsm.output, target)
        self.lsm.updateGradInput(input, self.nll.gradInput)
        self.gradInput = self.lsm.gradInput.view(size)
        return self.gradInput