TorchMetrics

Warnings

• breaking Python 3.9 support has been dropped with the release of v1.9.0. The minimum required Python version is now 3.10.
  Fix: Upgrade your Python environment to 3.10 or newer, or pin `torchmetrics<1.9.0`.
• breaking The default value for the `average` argument in `DiceScore` has changed from `None` to `"macro"` starting from v1.9.0. This can alter the behavior of existing code if the `average` argument was not explicitly set.
  Fix: Explicitly set the `average` argument in `DiceScore` to `None` or your desired reduction method if you relied on the previous default behavior.
• gotcha Metrics maintain internal states that accumulate data. Mixing these states across different phases (e.g., training, validation, testing) or re-using the same metric instance without resetting can lead to incorrect results or memory leaks.
  Fix: Always initialize separate metric instances for different phases (training, validation, test) or call `metric.reset()` after each complete evaluation epoch/phase to clear its internal state.
• gotcha Metric states are initialized on the CPU. When working with PyTorch tensors on GPU, especially in distributed training (DDP), ensure that metric objects are moved to the same device as the input data using `.to(device)`. Failure to do so can result in `RuntimeError: Encountered different devices in metric calculation`.
  Fix: Call `metric.to(device)` after initialization, or ensure the metric is registered as a child module within a `torch.nn.Module` or `LightningModule`, which handles device transfers automatically.
• gotcha When defining metrics as part of a `torch.nn.Module` or `LightningModule`, avoid using native Python `list` or `dict` to store multiple `Metric` instances. These will not be correctly identified as child modules, preventing automatic device placement and state management.
  Fix: Use `torch.nn.ModuleList` or `torch.nn.ModuleDict` instead of native Python collections when nesting metrics within a `torch.nn.Module`.
• gotcha Users of `MetricCollection` might encounter `UserWarning: The compute method of metric X was called before the update method...` This often indicates an issue where internal states of grouped metrics are not being updated correctly before `compute` is called, particularly in older versions or specific usage patterns.
  Fix: Ensure all metrics within the `MetricCollection` receive `update` calls for the relevant data. If issues persist, consider isolating metrics or upgrading to the latest `torchmetrics` version as device management and state handling are continually improved.
• gotcha For performance-critical applications, especially with `MeanMetric`, explicitly providing a weight tensor to `update` instead of relying on default values can be beneficial. Additionally, disabling NaN checks in the base `Aggregator` class or careful device management can reduce overhead.
  Fix: Profile your code. For `MeanMetric`, consider `metric.update(value, weight=my_tensor)`. For advanced optimization, explore `Aggregator` configurations and carefully manage device placement and cross-device synchronization.

Install

• pip install torchmetrics PyPI

Imports

• Accuracy

  from torchmetrics import Accuracy

• functional.accuracy

  from torchmetrics.functional import accuracy

• MetricCollection

  from torchmetrics import MetricCollection

• Metric

  from torchmetrics import Metric

  Base class for implementing custom metrics.

Quickstart

This quickstart demonstrates the core ways to use TorchMetrics: the functional API for stateless, single-batch computation, the class-based API for accumulating states over multiple batches, and MetricCollection for grouping several metrics. Remember to reset class-based metrics after each epoch or evaluation phase to avoid mixing states.

import torch
import torchmetrics
from torchmetrics import Accuracy, MetricCollection
from torchmetrics.functional import accuracy

# 1. Functional API: For single-batch, stateless computation
preds_f = torch.randn(10, 5).softmax(dim=-1)
target_f = torch.randint(5, (10,))
acc_functional = accuracy(preds_f, target_f, task="multiclass", num_classes=5)
print(f"Functional Accuracy: {acc_functional.item()}")

# 2. Class-based API: For accumulating metrics over multiple batches/epochs
metric = Accuracy(task="multiclass", num_classes=5)
preds_c = torch.randn(10, 5).softmax(dim=-1)
target_c = torch.randint(5, (10,))
metric.update(preds_c, target_c)

# Simulate another batch
preds_c2 = torch.randn(10, 5).softmax(dim=-1)
target_c2 = torch.randint(5, (10,))
metric.update(preds_c2, target_c2)

final_acc = metric.compute()
print(f"Class-based Accuracy (accumulated): {final_acc.item()}")
metric.reset() # Reset metric states for the next epoch/evaluation

# 3. MetricCollection: Group multiple metrics
metrics = MetricCollection({
    'Accuracy': Accuracy(task="multiclass", num_classes=5),
    'F1Score': torchmetrics.F1Score(task="multiclass", num_classes=5)
})
preds_mc = torch.randn(10, 5).softmax(dim=-1)
target_mc = torch.randint(5, (10,))
metrics.update(preds_mc, target_mc)
result_mc = metrics.compute()
print(f"MetricCollection Result: {result_mc}")