Convert a custom accuracy function into torchmetric?

[Abhinav43]

Hello, Thank you for this awesome project. I am new to use torchmetric, and I have one doubt regarding the custom function.

So if I have my own function to calculate the accuracy, such as :

From sklearn.metrics import f1_score

def get_hat(y_hat, threshold):
    hat_y = np.where(np.array(y_hat) > threshold, 1, 0)
    return hat_y

def get_acc(pred_y, y, threshold):

  if type(threshold) !=list:
    threshold = [threshold]
  
  final = {'acc': [], 'thres': []}
  for k in threshold:

    pred_ = get_hat(pred_y, k)
    real_ = y

    # some custom accuracy function defined somewhere else in file
    # just for example I am using sklearn f1-score

    acc   = f1_score(real_, pred_, average='macro')
    final['acc'].append(acc)
    final['thres'].append(k)

  return pd.DataFrame(final)

How I can convert this function into torch metric where all the internal logic will be from my function, Just want to wrap my logic in torch metric.

Thanks!

[SkafteNicki]

Hi @Abhinav43,
We have this page in the documentation: https://torchmetrics.readthedocs.io/en/latest/pages/implement.html which should help you get stated (I am assuming that you want to turn it into a modular metric). If this does not help you then feel free to reach out to me again.

[Borda]

@SkafteNicki pointed to the docs — let me give a concrete end-to-end example that should cover most "I have a custom function, how do I wrap it?" cases (v1.9.0).

Say you have this custom accuracy function:

def my_custom_accuracy(preds, target, top_k=3):
    """Top-k accuracy: correct if true label is in top-k predictions."""
    top_k_preds = preds.topk(top_k, dim=-1).indices
    correct = (top_k_preds == target.unsqueeze(-1)).any(dim=-1)
    return correct.sum(), target.shape[0]

Wrap it as a TorchMetrics Metric:

import torch
from torchmetrics import Metric


class MyTopKAccuracy(Metric):
    # Metadata — helps Lightning and logging tools
    is_differentiable = False
    higher_is_better = True
    full_state_update = False  # update() doesn't need full state

    def __init__(self, top_k: int = 3, **kwargs):
        super().__init__(**kwargs)
        self.top_k = top_k
        # Register states with reduction function for DDP
        self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
        correct, total = my_custom_accuracy(preds, target, self.top_k)
        self.correct += correct
        self.total += total

    def compute(self) -> torch.Tensor:
        return self.correct.float() / self.total

Usage:

metric = MyTopKAccuracy(top_k=5)

# Per-batch
for preds, target in dataloader:
    metric.update(preds, target)

# Epoch-level result (works across DDP automatically)
result = metric.compute()
metric.reset()

Key rules:

1. Always use self.add_state(...) — never plain self.foo = 0. Without add_state, reset/sync won't work.
2. Pick the right dist_reduce_fx: "sum" for counts, "cat" for lists of per-batch values, "mean" if each batch produces a final scalar.
3. Set full_state_update=False if your update() only appends/adds (doesn't read existing state). This enables a 2x speedup on forward().

Docs: Implementing Custom Metrics