TorchMetrics: 机器学习中强大而灵活的指标计算工具

TorchMetrics简介

TorchMetrics是一个专门用于机器学习指标计算的Python库,基于PyTorch开发。它提供了一套标准化的接口来实现各种常用的评估指标,具有以下主要特点:

提供100多种预定义的指标实现,涵盖分类、回归、检测等多个领域
支持自动批处理累积,无需手动管理状态
针对分布式训练进行了优化,可自动在多个设备间同步
提供简单的API用于创建自定义指标
与PyTorch Lightning无缝集成,进一步简化使用

通过使用TorchMetrics,研究人员和工程师可以大大减少在指标计算上的工作量,提高代码的可复用性和可维护性。

主要功能

模块化指标

TorchMetrics的核心是模块化指标(Module metrics)。这些指标被实现为PyTorch模块,内部维护计算所需的状态。主要优势包括:

自动在多个batch间累积结果
在分布式训练中自动同步多个设备的状态
支持指标间的算术运算

使用示例:

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)

    # 计算当前batch的指标
    acc = metric(preds, target)
    print(f"Batch {i} Accuracy: {acc}")

# 计算所有batch的累积指标  
acc = metric.compute()
print(f"Overall Accuracy: {acc}")

函数式指标

除了模块化指标,TorchMetrics还提供了函数式API。这些是简单的Python函数,接受张量作为输入并返回指标值。适用于不需要维护状态的场景。

使用示例:

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
)
print(f"Accuracy: {acc}")

自定义指标

TorchMetrics提供了简单的API用于创建自定义指标。只需继承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):
        preds = preds.argmax(dim=-1)
        assert preds.shape == target.shape
        self.correct += torch.sum(preds == target)
        self.total += target.numel()

    def compute(self):
        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))

可视化支持

TorchMetrics为大多数模块化指标内置了绘图支持。只需调用.plot()方法即可生成可视化结果:

import torch
from torchmetrics.classification import MulticlassAccuracy, MulticlassConfusionMatrix

num_classes = 3

# 生成模拟数据
w = torch.randn(num_classes)
target = lambda it: torch.multinomial((it * w).softmax(dim=-1), 100, replacement=True)
preds = lambda it: torch.multinomial((it * w).softmax(dim=-1), 100, replacement=True)

acc = MulticlassAccuracy(num_classes=num_classes, average="micro")
acc_per_class = MulticlassAccuracy(num_classes=num_classes, average=None)
confmat = MulticlassConfusionMatrix(num_classes=num_classes)

# 绘制单值指标
for i in range(5):
    acc_per_class.update(preds(i), target(i))
    confmat.update(preds(i), target(i))
fig1, ax1 = acc_per_class.plot()
fig2, ax2 = confmat.plot()

# 绘制多值指标
values = []
for i in range(10):
    values.append(acc(preds(i), target(i)))
fig3, ax3 = acc.plot(values)