Lightning-AI / Lightning-AI/torchmetrics
Parameter of MulticlassStatScores and MultilabelStatScores to control which classes/labels to include the averages
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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