Allow passing labels in (n_samples, n_classes) to AveragePrecision

[discort]

🚀 Feature

Allow passing labels in (n_samples, n_classes) to AveragePrecision like in sklearn.metrics average_precision_score

Motivation

ap = AveragePrecision(
    num_classes=7,
    pos_label=1,
    compute_on_step=False,
    dist_sync_on_step=False)
preds = torch.randn(8, 7)
target = torch.randint(0, 2, (8,7))
ap(preds, target)
ap.compute()
~/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py in precision_recall_curve(preds, target, num_classes, pos_label, sample_weights)
    248         [tensor([0.7500]), tensor([0.7500]), tensor([0.0500, 0.7500]), tensor([0.0500, 0.7500]), tensor([0.0500])]
    249     """
--> 250     preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)
    251     return _precision_recall_curve_compute(preds, target, num_classes, pos_label, sample_weights)

~/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py in _precision_recall_curve_update(preds, target, num_classes, pos_label)
    107
    108     else:
--> 109         raise ValueError("preds and target must have same number of dimensions, or one additional dimension for preds")
    110
    111     return preds, target, num_classes, pos_label

ValueError: preds and target must have same number of dimensions, or one additional dimension for preds

Pitch

In sklearn there is ability to calculate average precision for multilabel target:

from sklearn.metrics import average_precision_score
import numpy as np

target = np.random.randint(0, 2, (8, 7))
preds = np.random.randn(8, 7)

average_precision_score(target , preds)
0.505187074829932

Alternatives

Additional context

Version:

torchmetrics==0.4.1

[github-actions]

Hi! thanks for your contribution!, great first issue!

[discort]

cc @Borda @SkafteNicki

[SkafteNicki]

Hi @discort thanks for your issue.
Please feel free to send a pull request with the enchancement :]

[discort]

thanks for getting back to me @SkafteNicki

I found that AveragePrecision is used to work with multilabel targets in pytorch-lightning==1.2.8

def main():
    import torch
    from pytorch_lightning.metrics import AveragePrecision
    ap = AveragePrecision(
        num_classes=7,
        pos_label=1,
        compute_on_step=False,
        dist_sync_on_step=False)
    preds = torch.randn(8, 7)
    target = torch.randint(0, 2, (8,7))
    ap(preds, target)
    res = ap.compute()
    print('res=', res)


if __name__ == "__main__":
    main()

output:

> python test_avg_precision.py
/Users/discort/git/package/.env/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:68: UserWarning: Metric `AveragePrecision` will save all targets and predictions in buffer. For large datasets this may lead to large memory footprint.
  warnings.warn(*args, **kwargs)
res= tensor(0.5896)

but in pytorch-lightning==1.3.8 and torchmetrics==0.4.1 it stopped working, so it seems this is a bug:

def main():
    import torch
    from torchmetrics import AveragePrecision
    ap = AveragePrecision(
        num_classes=7,
        pos_label=1,
        compute_on_step=False,
        dist_sync_on_step=False)
    preds = torch.randn(8, 7)
    target = torch.randint(0, 2, (8,7))
    ap(preds, target)
    res = ap.compute()
    print('res=', res)


if __name__ == "__main__":
    main()

output:

> python test_avg_precision.py
/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/utilities/prints.py:37: UserWarning: Metric `AveragePrecision` will save all targets and predictions in buffer. For large datasets this may lead to large memory footprint.
  warnings.warn(*args, **kwargs)
Traceback (most recent call last):
  File "test_avg_precision.py", line 143, in <module>
    main()
  File "test_avg_precision.py", line 138, in main
    res = ap.compute()
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/metric.py", line 370, in wrapped_func
    self._computed = compute(*args, **kwargs)
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/classification/average_precision.py", line 132, in compute
    return _average_precision_compute(preds, target, self.num_classes, self.pos_label)
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/average_precision.py", line 42, in _average_precision_compute
    precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py", line 177, in _precision_recall_curve_compute
    return _precision_recall_curve_compute_multi_class(preds, target, num_classes, sample_weights)
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py", line 151, in _precision_recall_curve_compute_multi_class
    res = precision_recall_curve(
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py", line 250, in precision_recall_curve
    preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)
  File "/Users/discort/git/package/.env/lib/python3.8/site-packages/torchmetrics/functional/classification/precision_recall_curve.py", line 109, in _precision_recall_curve_update
    raise ValueError("preds and target must have same number of dimensions, or one additional dimension for preds")
ValueError: preds and target must have same number of dimensions, or one additional dimension for preds

I can check what changed if you are unable how to solve it quickly.