Lightning-AI / Lightning-AI/torchmetrics

Parameter of MulticlassStatScores and MultilabelStatScores to control which classes/labels to include the averages

Open
#1,723 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
2.5k
Forks
526
Avg merge
6d 11h
Merged PRs (30d)
5

Description

## 🚀 Feature

Add a parameter to `MulticlassStatScores` and `MultilabelStatScores` to control which classes/labels to include in averaging.

### Motivation

Sklearn's [`precision_recall_fscore_support`](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.precision_recall_fscore_support.html#sklearn.metrics.precision_recall_fscore_support) allows users to define the labels used in averaging the computed metrics (as well as the order if the metrics are not averaged). This allows calculating "a multiclass average ignoring a majority negative class". E.g. in my use-case (sequence tagging), I do want to consider datapoints which have an out-tag ("O", meaning they are not tagged), as they might contribute to the false positives of other classes. Hence, `ignore_index` is not sufficient, as the datapoints would be completely excluded.

### Pitch

Add a parameter `classes` to `MulticlassStatScores` and `labels` to `MultilabelStatScores` to limit the calculation of true positives, fp, fn, and tn to these classes/labels. The resulting averages (e.g. f1-score, accuracy) would then be an average only of the selected classes/labels.

If the `classes` / `labels` parameter is specified, `num_classes` / `num_labels` would not need to be set (or if they are set and do not agree with the passed number of classes/labels, an Exception would need to be raised).

### Alternatives

Currently, I am using a very hacky solution:

```python
def metric_with_certain_labels_only(
metric_type: Union[Type[MulticlassStatScores], Type[MultilabelStatScores]],
included_labels: torch.Tensor,
average: Optional[Literal["micro", "macro", "weighted", "none"]] = "macro",
**kwargs,
):
metric_average = average

if average == "micro":
metric_average = "none"

metric = metric_type(average=metric_average, **kwargs)

_final_state_inner = metric._final_state

def _final_state_wrapper():
state = _final_state_inner()

# manipulate the state variable
new_state = (s[torch.tensor(included_labels)] for s in state)

return new_state

metric._final_state = _final_state_wrapper # type: ignore

if average == "micro":
compute_inner = metric.compute

def compute_wrapper():
metric.average = "micro"
result = compute_inner()
metric.average = "none"
return result

metric.compute = compute_wrapper # type: ignore

return metric
```

I am not happy with this solution for two reasons:

1. It would be nicer, if included_labels would be part of the metric init signature directly or (at least) would be provided in a wrapper. It is possible to rewrite this helper function as a Wrapper, however this would require changing the averaging on an already created metric (and recreating an new state), and (more importantly)
2. since the classes/labels are selected on the read-out of the state (which requires fewer changes), relevant stats need to be tracked for all classes (even the once which are not included). This is especially inefficient in case `average` is set to `"micro"`.

Ideally, the stats would be reduced to the selected classes already in `_multiclass_stat_scores_update` / `_multilabel_stat_scores_update`

### Additional context

I had already opened a [discussion](https://github.com/Lightning-AI/torchmetrics/discussions/1721). However, I believe this cannot be solved without a new feature.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading MulticlassStatScores, MultilabelStatScores, and the _multiclass_stat_scores_update and _multilabel_stat_scores_update entry points. Trace how averaging and num_classes/num_labels are handled, then define coverage for selecting classes or labels while preserving the requested metric behavior and avoiding unnecessary state tracking.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.