Source code for catalyst.tools.meters.apmeter

"""
The APMeter measures the average precision per class.
"""
import math

import torch

from catalyst.tools.meters import meter


class APMeter(meter.Meter):
    """
    The APMeter is designed to operate on `NxK` Tensors `output` and
    `target`, and optionally a `Nx1` Tensor weight where:

    1. The `output` contains model output scores for `N` examples and
    `K` classes that ought to be higher when the model is more convinced
    that the example should be positively labeled, and smaller when the
    model believes the example should be negatively labeled
    (for instance, the output of a sigmoid function).

    2. The `target` contains only values 0 (for negative examples)
    and 1 (for positive examples).

    3. The `weight` ( > 0) represents weight
    for each sample.
    """

    def __init__(self):
        """Constructor method for the ``APMeter`` class."""
        super(APMeter, self).__init__()
        self.reset()

    def reset(self):
        """Resets the meter with empty member variables."""
        self.scores = torch.FloatTensor(torch.FloatStorage())
        self.targets = torch.LongTensor(torch.LongStorage())
        self.weights = torch.FloatTensor(torch.FloatStorage())

    def add(
        self,
        output: torch.Tensor,
        target: torch.Tensor,
        weight: torch.Tensor = None,
    ) -> None:
        """Add a new observation.

        Args:
            output (Tensor): NxK tensor that for each of the N examples
                indicates the probability of the example belonging to each of
                the K classes, according to the model. The probabilities should
                sum to one over all classes
            target (Tensor): binary NxK tensort that encodes which of the K
                classes are associated with the N-th input
                (eg: a row [0, 1, 0, 1] indicates that the example is
                associated with classes 2 and 4)
            weight (optional, Tensor): Nx1 tensor representing the weight for
                each example (each weight > 0)
        """
        if not torch.is_tensor(output):
            output = torch.from_numpy(output)
        if not torch.is_tensor(target):
            target = torch.from_numpy(target)

        if weight is not None:
            if not torch.is_tensor(weight):
                weight = torch.from_numpy(weight)
            weight = weight.squeeze()
        if output.dim() == 1:
            output = output.view(-1, 1)
        else:
            assert output.dim() == 2, (
                "wrong output size "
                "(should be 1D or 2D with one column per class)"
            )

        if target.dim() == 1:
            target = target.view(-1, 1)
        else:
            assert target.dim() == 2, (
                "wrong target size "
                "(should be 1D or 2D with one column per class)"
            )

        if weight is not None:
            assert weight.dim() == 1, "Weight dimension should be 1"
            assert weight.numel() == target.size(
                0
            ), "Weight dimension 1 should be the same as that of target"
            assert torch.min(weight) >= 0, "Weight should be non-negative only"
        assert torch.equal(
            target ** 2, target
        ), "targets should be binary (0 or 1)"
        if self.scores.numel() > 0:
            assert target.size(1) == self.targets.size(
                1
            ), "dimensions for output should match previously added examples."

        # make sure storage is of sufficient size
        if self.scores.storage().size() < self.scores.numel() + output.numel():
            new_size = math.ceil(self.scores.storage().size() * 1.5)
            new_weight_size = math.ceil(self.weights.storage().size() * 1.5)
            self.scores.storage().resize_(int(new_size + output.numel()))
            self.targets.storage().resize_(int(new_size + output.numel()))
            if weight is not None:
                self.weights.storage().resize_(
                    int(new_weight_size + output.size(0))
                )

        # store scores and targets
        offset = self.scores.size(0) if self.scores.dim() > 0 else 0
        self.scores.resize_(offset + output.size(0), output.size(1))
        self.targets.resize_(offset + target.size(0), target.size(1))
        self.scores.narrow(0, offset, output.size(0)).copy_(output)
        self.targets.narrow(0, offset, target.size(0)).copy_(target)

        if weight is not None:
            self.weights.resize_(offset + weight.size(0))
            self.weights.narrow(0, offset, weight.size(0)).copy_(weight)

    def value(self) -> torch.Tensor:
        """Returns the model"s average precision for each class.

        Returns:
            torch.Tensor: 1xK tensor, with avg precision for each class k
        """
        if self.scores.numel() == 0:
            return 0
        ap = torch.zeros(self.scores.size(1))
        if hasattr(torch, "arange"):
            rg = torch.arange(1, self.scores.size(0) + 1).float()
        else:
            rg = torch.range(1, self.scores.size(0)).float()
        if self.weights.numel() > 0:
            weight = self.weights.new(self.weights.size())
            weighted_truth = self.weights.new(self.weights.size())

        # compute average precision for each class
        for k in range(self.scores.size(1)):
            # sort scores
            scores = self.scores[:, k]
            targets = self.targets[:, k]
            _, sortind = torch.sort(scores, dim=0, descending=True)
            truth = targets[sortind]
            if self.weights.numel() > 0:
                weight = self.weights[sortind]
                weighted_truth = truth.float() * weight
                rg = weight.cumsum(0)

            # compute true positive sums
            if self.weights.numel() > 0:
                tp = weighted_truth.cumsum(0)
            else:
                tp = truth.float().cumsum(0)

            # compute precision curve
            precision = tp.div(rg)

            # compute average precision
            ap[k] = precision[truth.byte()].sum() / max(float(truth.sum()), 1)
        return ap