Repository: Lightning-AI/torchmetrics

Issue:
Allow passing labels in (n_samples, n_classes) to AveragePrecision
## 🚀 Feature
Allow passing labels in `(n_samples, n_classes)` to [AveragePrecision](https://torchmetrics.readthedocs.io/en/latest/references/modules.html#averageprecision) like in sklearn.metrics [average_precision_score](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.average_precision_score.html)

### 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

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### Additional context
Version:
```
torchmetrics==0.4.1
```

Allow passing labels in (n_samples, n_classes) to AveragePrecision
## 🚀 Feature
Allow passing labels in `(n_samples, n_classes)` to [AveragePrecision](https://torchmetrics.readthedocs.io/en/latest/references/modules.html#averageprecision) like in sklearn.metrics [average_precision_score](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.average_precision_score.html)

### 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

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### Additional context
Version:
```
torchmetrics==0.4.1
```



Discussion:
Hi! thanks for your contribution!, great first issue!
cc @Borda @SkafteNicki 
Hi @discort thanks for your issue.
Please feel free to send a pull request with the enchancement :]
thanks for getting back to me @SkafteNicki 

I found that `AveragePrecision` is used to work with multilabel targets in `pytorch-lightning==1.2.8`
```py
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:
```py
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.
Hi! thanks for your contribution!, great first issue!
cc @Borda @SkafteNicki 
Hi @discort thanks for your issue.
Please feel free to send a pull request with the enchancement :]
thanks for getting back to me @SkafteNicki 

I found that `AveragePrecision` is used to work with multilabel targets in `pytorch-lightning==1.2.8`
```py
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:
```py
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.

Resolving patch:
diff --git a/torchmetrics/functional/classification/precision_recall_curve.py b/torchmetrics/functional/classification/precision_recall_curve.py
--- a/torchmetrics/functional/classification/precision_recall_curve.py
+++ b/torchmetrics/functional/classification/precision_recall_curve.py
@@ -148,13 +148,20 @@ def _precision_recall_curve_compute_multi_class(
     precision, recall, thresholds = [], [], []
     for cls in range(num_classes):
         preds_cls = preds[:, cls]
-        res = precision_recall_curve(
+
+        prc_args = dict(
             preds=preds_cls,
             target=target,
             num_classes=1,
             pos_label=cls,
             sample_weights=sample_weights,
         )
+        if target.ndim > 1:
+            prc_args.update(dict(
+                target=target[:, cls],
+                pos_label=1,
+            ))
+        res = precision_recall_curve(**prc_args)
         precision.append(res[0])
         recall.append(res[1])
         thresholds.append(res[2])


Test patch:
diff --git a/tests/classification/test_average_precision.py b/tests/classification/test_average_precision.py
--- a/tests/classification/test_average_precision.py
+++ b/tests/classification/test_average_precision.py
@@ -21,6 +21,7 @@
 from tests.classification.inputs import _input_binary_prob
 from tests.classification.inputs import _input_multiclass_prob as _input_mcls_prob
 from tests.classification.inputs import _input_multidim_multiclass_prob as _input_mdmc_prob
+from tests.classification.inputs import _input_multilabel
 from tests.helpers import seed_all
 from tests.helpers.testers import NUM_CLASSES, MetricTester
 from torchmetrics.classification.average_precision import AveragePrecision
@@ -55,6 +56,12 @@ def _sk_avg_prec_multiclass_prob(preds, target, num_classes=1):
     return _sk_average_precision_score(y_true=sk_target, probas_pred=sk_preds, num_classes=num_classes)
 
 
+def _sk_avg_prec_multilabel_prob(preds, target, num_classes):
+    sk_preds = preds.reshape(-1, num_classes).numpy()
+    sk_target = target.view(-1, num_classes).numpy()
+    return sk_average_precision_score(sk_target, sk_preds, average=None)
+
+
 def _sk_avg_prec_multidim_multiclass_prob(preds, target, num_classes=1):
     sk_preds = preds.transpose(0, 1).reshape(num_classes, -1).transpose(0, 1).numpy()
     sk_target = target.view(-1).numpy()
@@ -66,6 +73,7 @@ def _sk_avg_prec_multidim_multiclass_prob(preds, target, num_classes=1):
         (_input_binary_prob.preds, _input_binary_prob.target, _sk_avg_prec_binary_prob, 1),
         (_input_mcls_prob.preds, _input_mcls_prob.target, _sk_avg_prec_multiclass_prob, NUM_CLASSES),
         (_input_mdmc_prob.preds, _input_mdmc_prob.target, _sk_avg_prec_multidim_multiclass_prob, NUM_CLASSES),
+        (_input_multilabel.preds, _input_multilabel.target, _sk_avg_prec_multilabel_prob, NUM_CLASSES),
     ]
 )
 class TestAveragePrecision(MetricTester):