torchmetrics

TorchMetrics项目介绍

简介

TorchMetrics是一个专为PyTorch应用设计的、高效且可扩展的机器学习指标库。它包括超过100种的内置指标，同时提供了一套易于使用的API，方便用户创建自定义指标。其主要特点包括：

提供标准化接口，提高了模型结果的可重复性
减少样板代码，简化开发流程
自动累积批处理数据，优化分布式训练
在多个设备间自动同步数据

用户可以在任何PyTorch模型中使用TorchMetrics，也可以结合PyTorch Lightning一起使用，以便享受额外的功能，如自动将模块化指标放置在正确的设备上并支持记录指标。

安装

TorchMetrics的安装非常简单。可以通过以下几种方法进行安装：

使用PyPI安装：
```
pip install torchmetrics
```

使用Conda进行安装：

conda install -c conda-forge torchmetrics

从源码安装：

pip install git+https://github.com/Lightning-AI/torchmetrics.git@release/stable

安装额外依赖以支持不同领域的指标（如音频、图像、文本等）：

pip install torchmetrics[audio]
pip install torchmetrics[image]
pip install torchmetrics[text]
pip install torchmetrics[all]  # 下载安装所有依赖

使用指南

模块化指标

模块化指标利用内部的指标状态（类似于PyTorch模块的参数），可以在多个设备中自动进行累积和同步。以下是示例代码，展示了如何在单个GPU或CPU上使用这些指标：

import torch
import torchmetrics

# 初始化指标
metric = torchmetrics.classification.Accuracy(task="multiclass", num_classes=5)

# 将指标放置到相应的设备
device = "cuda" if torch.cuda.is_available() else "cpu"
metric.to(device)

n_batches = 10
for i in range(n_batches):
    # 模拟分类问题
    preds = torch.randn(10, 5).softmax(dim=-1).to(device)
    target = torch.randint(5, (10,)).to(device)

    # 当前批次的指标
    acc = metric(preds, target)
    print(f"第 {i} 批次的准确率: {acc}")

# 计算所有批次的总指标
acc = metric.compute()
print(f"所有数据的准确率: {acc}")

自定义指标

用户可以通过继承torchmetrics.Metric来实现自己的指标。只需实现update和compute方法即可：

import torch
from torchmetrics import Metric

class MyAccuracy(Metric):
    def __init__(self):
        super().__init__()
        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:
        preds = preds.argmax(dim=-1)
        assert preds.shape == target.shape
        self.correct += torch.sum(preds == target)
        self.total += target.numel()

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

my_metric = MyAccuracy()
preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))

print(my_metric(preds, target))

功能性指标

除了模块化指标，每个指标也有其功能性版本，这些版本是简单的Python函数，输入torch.tensors，输出相应的指标结果。

import torch
import torchmetrics

preds = torch.randn(10, 5).softmax(dim=-1)
target = torch.randint(5, (10,))

acc = torchmetrics.functional.classification.multiclass_accuracy(
    preds, target, num_classes=5
)