torchmetrics指标状态管理:从单卡累加到多卡DDP同步实践
1. 从"能累加的指标"到"不能累加的指标":手写统计失效的分界线
第一次在训练脚本里删掉自己写的correct += (pred == target).sum(),换成 torchmetrics 的Accuracy,其实是被迫的。前面几个 epoch 都相安无事,直到换成 AUROC、F1 这种指标,才发现累加变量这条路根本走不通。torchmetrics 是 PyTorch 生态里专门做指标计算与状态管理的库,它管的不只是"算一个数",而是把这些数在整个训练、验证、多卡、跨 epoch 的过程中怎么存、怎么合、什么时候清空这一整套脏活全接走。如果你正在写训练循环,尤其是验证集上要报三五个指标,或者手上有多张卡要跑 DDP,那么它值得花半小时认真看一遍;如果你只是偶尔算个准确率,用torchmetrics.functional里的无状态函数就够了,不必引入对象的复杂度。
1.1 那个在验证集上出过问题的累加写法
常见的写法大概是这样:初始化correct = 0、total = 0,每个 batch 执行correct += (pred.argmax(1) == target).sum().item(),最后acc = correct / total。这个写法本身没错,因为它累加的是"正确样本数"和"样本总数"这两个可以直接相加的量,比值自然就对了。问题出在两个地方:一是在验证循环结束后忘了把这两个变量清零,下个 epoch 接着往上加,指标会随训练进程慢慢飘向 1.0;二是当你想多加一个 F1 或者 AUC 时,就会突然发现没有一个"正确的个数"可以对着加,代码开始变得别扭,于是有人退回用sklearn.metrics在 CPU 上把整个验证集预测结果存下来最后统一算,显存和内存一起被拖着走。
sklearn的路线在验证集不大时完全可行,真出问题是在训练集上也要监控指标的时候,或者数据加载器一秒钟吐几十个 batch、你想要一个实时的滑动指标看趋势。每算一次就把几千个预测搬回 CPU 做一次 numpy 转换,开销会以你意想不到的方式出现在 profile 里。
1.2 平均值再平均,为什么在最后一批样本上翻车
比忘记清零更隐蔽的是"先算每个 batch 的指标,再对所有 batch 的指标求平均"。这个写法在 batch 大小完全一致时会得到和真实值很接近的结果,但只是接近,不是相等——因为每个 batch 内部的比例是四舍五入意义上的近似。真正翻车的是最后一个 batch:当数据集大小不能被 batch_size 整除时,drop_last=False会让最后一批的样本数少很多。举个例子,前 9 个 batch 每个 100 条,正确率都是 80%,最后一批只有 10 条,正确率 20%。直接对 10 个数字求平均得到 74%,而真实准确率是 (720 + 2) / 910 ≈ 79.3%。差了五个点,足以让你误判一次调参是否有效。
这个错误在类别数多、单 batch 内样本分布不均的场景下会被进一步放大。AI 圈子里有句话叫"指标骗人",大部分时候骗人的不是指标本身,而是平均的方式。所以判断一个指标能不能"先算后平均",关键看它的定义是比率型还是可分解型:准确率、召回率、精确率分母是样本数,必须累加分子分母;而像每个 batch 的 loss 求平均,只要 batch 大小一致就没问题,不一致同样要用加权平均。
1.3 torchmetrics 换掉的其实不是函数,是状态
理解 torchmetrics 的关键,是别把它当成一个"计算函数库"。torchmetrics.functional那一层确实是一堆纯函数,输入张量输出标量,用完即弃;但torchmetrics.Metric那一层是有状态的对象,它在内部维护若干张量当作累加器,每次update()只往里加,compute()才把累加器合成为最终结果,reset()负责清零。这个三段式设计看起来只是把变量封装进类里,实际带来的好处是:多卡场景下它能自动在compute()时跨进程同步这些中间量;state_dict能把半途的状态存盘再恢复;MetricCollection能把十几个指标打包成一次update调用,省掉重复的argmax和softmax。
换句话说,你交给它的不是一个 batch 的预测,而是一份"还没算完的中间结果",它替你把这份结果在整个生命周期里管好。
2. Metric 的状态机:update、compute、reset 各自的真实职责
写自定义指标之前,先把这三个方法的边界搞清楚,能省掉大量调试时间。很多人第一次继承Metric时会把计算逻辑写在update里,每次 update 就算一遍准确率存进一个 list,然后compute里对 list 求平均——这就又回到了上一节说的错误平均法。正确的心智模型是:update只做累加,不做除法;compute只做合成,不做累加;reset只做归零,不做计算。
2.1 add_state 声明的到底是什么
在__init__里调用add_state(name, default, dist_reduce_fx, persistent),本质上是给这个 Metric 注册一个可被框架管理的张量槽位。default决定初始值和形状,dist_reduce_fx决定多卡同步时怎么合并("sum"、"mean"、"cat"、"min"、"max"或者一个自定义函数),persistent决定它要不要进state_dict。这三个参数里最容易出错的是dist_reduce_fx,因为它默认是None,也就是不做任何跨进程合并——单卡跑没问题,一上多卡,每个进程算出自己那份结果,日志里就会出现 N 个不同的数字,你可能以为是随机种子的问题,排查半天。
add_state注册的东西会被reset()恢复成default的副本,会被.to(device)一起搬到目标设备,也会参与state_dict的序列化。这意味着你不需要手动写if self.tp is None: self.tp = torch.zeros(...)这类惰性初始化代码,框架已经管了。
提示:
default传入的如果是张量,add_state存的是它的克隆,多个 Metric 实例之间不会共享同一块内存。但如果你传的是可变对象比如列表,务必确认框架版本的处理方式,老版本在 reset 时可能出现浅拷贝导致的状态残留。
2.2 compute 的返回值为什么会被"记住"
这是很多人踩过的坑:在验证循环里先调用一次compute()打印结果,然后继续update()几个 batch,再compute(),发现数字没变。原因是Metric内部有一个_computed标志和一份结果缓存,compute()只有在状态被标记为"脏"(即发生过update或reset)时才会真正重算,否则直接返回上次的结果。这个设计是为了性能——日志系统经常会对同一个指标多次取值,没必要重复计算。
要规避这个问题,要么每次compute()之后立刻reset(),要么保证一个完整的评估阶段内只compute()一次。如果你的日志框架(比如各种实验记录工具)会在同一个 step 里多次读取指标值,缓存反而是好事,它保证了同一个 step 打印出来的数字是一致的。
2.3 forward 和 compute 不是一回事
metric(preds, target)这种直接调用的写法,走的是forward而不是compute。默认情况下(full_state_update=True),forward的内部动作是:先把当前累积的状态备份下来,reset清空,只用这一批数据做一次update加compute,拿到结果,再把备份的状态恢复回去。所以forward返回的是当前这一批数据单独的指标值,跟你之前累积的历史无关。
这个语义差异非常关键。训练时想看"这个 batch 表现如何",用metric(preds, target);验证时想看"整个验证集表现如何",必须用update累积再compute。两者混用会出现"训练日志里的准确率忽高忽低、验证日志里纹丝不动"这种诡异现象。另外forward因为有状态备份和恢复的动作,开销比update大不少,在数据量大、batch 数多的循环里频繁调用会拖慢训练。
3. 分类指标选型的判断依据:从 Accuracy 到 AUROC
torchmetrics 的分类指标从 0.11 版本起统一成了task参数风格,MulticlassAccuracy、MultilabelF1Score、BinaryAUROC这样的具体类名取代了早期那个靠num_classes和average猜任务的Accuracy。老代码迁移时最容易卡住的就是这一层,因为参数名变了,行为也跟着变了。
3.1 输入形态决定你能用哪个指标
同样是准确率,喂进去的张量形状不同,走的路径完全不同。传[N]的整型类别索引,指标直接和目标比对;传[N, C]的浮点张量,指标内部先做argmax再比对。如果你的模型输出的是 logits,直接传给 Accuracy 是安全的,因为它内部会判断输入是不是在[0, 1]之间以及每行是否归一化,然后决定要不要做argmax。但如果模型输出的是已经softmax过的概率,也照样能传,结果一致。
真正需要区分的是 AUROC 这一类。它们需要的是每个类别的置信度分数,然后对每个类别单独扫阈值算曲线下面积。传[N]的硬标签进去,会直接报错。
| 指标 | 期望输入 | 是否接受 logits | 典型用途 |
|---|---|---|---|
| Accuracy | [N]索引 或[N, C]分数 | 接受 | 类别均衡的粗粒度评估 |
| F1Score | 同上 | 接受 | 类别不均衡 |
| AUROC | [N, C]概率,值域[0, 1] | 不接受 | 排序质量、阈值无关评估 |
| AveragePrecision | [N, C]概率 | 不接受 | 正负样本极不均衡 |
| CohenKappa | [N]索引 | 不接受 | 标注一致性 |
3.2 AUROC 为什么拒绝 logits
有个直觉上很合理的误解:二分类下 sigmoid 是单调递增函数,所以对 logits 排序和对概率排序结果一样,AUROC 应该不变。这个推理在二分类成立,但 torchmetrics 仍然会拦,因为它要做统一的输入校验。多分类场景下这个推理就崩了——softmax 对单个 logit 并不是单调的,分母里包含了其他类别的指数和,某个类别的 logit 升高时,如果其他类别升得更多,它的概率反而会降。所以多分类 AUROC 必须先过 softmax。
实测下来,MulticlassAUROC在收到超出[0, 1]范围的输入时会抛ValueError,报错信息里会明确让你检查是不是忘了做 softmax。这个报错其实是个好事,比悄悄算出一个错误的数字强。
import torch from torchmetrics.classification import MulticlassAUROC auroc = MulticlassAUROC(num_classes=10, average="macro").to(device) for images, targets in val_loader: images, targets = images.to(device), targets.to(device) logits = model(images) probs = torch.softmax(logits, dim=-1) auroc.update(probs, targets) print(auroc.compute()) auroc.reset()3.3 average 参数在小类别上的连锁反应
average="micro"把所有类别的预测摊平成一个二分类问题算,得到的数值通常接近整体准确率,看起来漂亮;average="macro"对每个类别单独算再取算术平均,小类别会被赋予和大类别相同的权重,数值往往低一截;average="weighted"按类别样本数加权,介于两者之间。选哪个不是审美问题,而是你要回答什么业务问题——想知道"模型整体对了多少",用 micro;想知道"模型对每个类别是否都还行",用 macro,它会把长尾类别的糟糕表现直接暴露在数字里。
我见过不止一次有人用 micro 的 F1 报了个 0.92,看起来很棒,换成 macro 掉到 0.61,一查发现有个类别只有几十条样本、全部预测错误。这不是指标的问题,是指标选择掩盖了问题。
4. 多卡训练下的状态同步:那些让指标对不齐的细节
单卡跑通的代码搬到 DDP 上,最常出现的现象是:每个进程打印出来的指标都不一样,而且都比你预期的低。这不是模型的问题,是每个进程只看到了自己那份数据。torchmetrics 的解决方案是在compute()时自动把这些中间状态跨进程合并,但能不能生效取决于你在add_state时给对了dist_reduce_fx。
4.1 默认同步发生在 compute,而不是 update
这个设计是刻意的。update()频繁调用,每个 batch 都同步一次的话通信开销会压垮训练吞吐;compute()一个 epoch 才调一次,同步一次的代价可以忽略。所以默认行为是:各进程独立累积本地状态,compute()时做一次 all-gather 或 all-reduce,合并后计算,结果在每个进程上一致。
需要留意的坑是:如果你的状态里存的是"平均值"而不是"总和",用dist_reduce_fx="mean"合并,在各进程样本数相同的情况下没问题;样本数不同时(比如最后一个 batch 分配不均,或者用了加权采样),直接对平均值求平均是错的,必须先合并总和与计数,最后再除。这就是为什么自定义指标时推荐把中间状态设计成"可加的量"。
4.2 dist_sync_on_step 打开之后发生了什么
dist_sync_on_step=True会让forward每次调用都触发同步。它解决的是"想在训练中途看实时全局指标"的需求,代价是每个 batch 一次通信。绝大多数情况下不需要打开它——你真要实时看趋势,也应该在固定间隔才记录一次,而不是每个 step 都同步。
还有一种情况是同步了但结果还是不对:指标对象没有.to(device)。如果模型在 GPU 上、指标对象还留在 CPU,update时框架会把输入张量搬到指标所在的设备,这个隐式拷贝不会报错,但会悄悄拖慢速度,并且当你的状态张量很大(比如保存全部预测用于算 AUROC)时,CPU 内存会被吃满。
4.3 sync_context 与 compute_on_cpu 的取舍
sync_context()是一个上下文管理器,在它内部调用update不会触发任何跨进程同步,退出时才统一同步一次。适合的场景是你在一个循环里做多次局部计算,或者用了多进程数据并行但想精细控制同步时机。
with metric.sync_context(): for batch in loader: metric.update(logits, targets) # 退出上下文后状态已经同步 result = metric.compute()另一个值得知道的参数是compute_on_cpu。像 AUROC、AveragePrecision 这类指标需要在compute时对全部累积的预测排序,如果累积了几万个[N, C]的浮点张量放在显存里,很容易 OOM。把它设成True会让状态张量留在 CPU 上,compute时在 CPU 算,慢一点但省显存。验证集大的时候这个开关能救命。
注意:
compute_on_cpu=True之后,update传入的 GPU 张量会被搬到 CPU,这个拷贝本身是有成本的。如果指标状态很小(比如只有几个计数器),没必要开;只有状态里的张量随样本数增长时才值得。
5. 自己写一个 Metric:把业务口径钉进框架
现成指标覆盖不了的场景比想象中多。比如"只统计置信度高于 0.8 的样本上的准确率"、"按类别加权但权重来自业务方的成本矩阵"、"把连续值分桶后算一致性"。这些用sklearn也能写,但你就失去了多卡同步和状态管理。自己继承Metric其实不难,难的是把状态设计对。
5.1init里必须做的三件事
第一件是super().__init__(),把dist_sync_on_step、compute_on_cpu这些参数透传上去;第二件是用add_state声明所有需要在update之间保留的张量,不要用普通属性存,否则.to(device)和state_dict都不会管它;第三件是把不随训练变化的常量(比如类别数、权重向量)用register_buffer注册成 buffer,这样它们会跟着.to(device)一起走。
import torch from torchmetrics import Metric class WeightedRecall(Metric): def __init__(self, num_classes: int, weights=None, **kwargs): super().__init__(**kwargs) self.num_classes = num_classes if weights is None: weights = torch.ones(num_classes) self.register_buffer("weights", weights.float()) self.add_state("tp", default=torch.zeros(num_classes), dist_reduce_fx="sum") self.add_state("pos", default=torch.zeros(num_classes), dist_reduce_fx="sum") def update(self, preds: torch.Tensor, target: torch.Tensor): if preds.ndim > 1: preds = preds.argmax(dim=-1) for c in range(self.num_classes): mask = target == c self.pos[c] += mask.sum() self.tp[c] += ((preds == c) & mask).sum() def compute(self) -> torch.Tensor: recall = self.tp / self.pos.clamp(min=1) w = self.weights.to(recall.device) return (recall * w).sum() / w.sum()这段代码把"加权召回率"拆成了两个可加的计数器,dist_reduce_fx="sum"保证多卡合并时分子分母分别相加,比值自然正确。如果图省事直接把每个 batch 的 recall 存起来用"mean"合并,样本数不均匀时结果就是错的。
5.2 reduce_fn 选错会让多卡结果翻倍
一个很典型的现象:单卡跑出来准确率 0.85,双卡跑出来 1.70。原因通常是状态初始化成了torch.zeros(1)但update里做了+=,而dist_reduce_fx设成了"cat"或者别的合并方式,导致 1.0 + 0.7 这种荒谬的加和。排查这类问题的第一步永远是:打印每个进程compute()之前的状态值,看看是不是数值直接相加了。
另一个隐蔽的坑是状态用float32累积但输入是float16。+=的时候 PyTorch 会做类型提升,一般不会出错,但如果你的状态张量在初始化时用了torch.zeros(n, dtype=torch.half),在样本量很大时累加会溢出成inf。稳妥做法是让所有计数器统一用float32或int64。
5.3 状态设计的一条经验法则
判断一个量该不该作为独立状态,问自己:它能不能在不丢失信息的前提下和其他进程的结果直接相加或拼起来?能,就用"sum"或"cat"存原始量;不能,就想办法拆成能相加的原始量。平均值拆成总和加计数,方差拆成平方和加计数,中位数这种拆不出来的就只能用"cat"把全部数据存下来最后统一算——这也意味着中位数类指标的内存会随样本量线性增长,要配合compute_on_cpu和定期 reset 使用。
6. 实际跑起来才暴露的六个问题
6.1 设备、dtype 与空 batch
指标对象没.to(device)是最常见的一个,前面提过,表现是不报错但变慢。dtype 问题更隐蔽:模型用了混合精度,输出是float16,传给 AUROC 时可能因为精度损失导致某些概率变成 0 或 1,排序出现大量并列,最终数值偏移。稳妥做法是在 update 之前显式.float()。空 batch 则会在compute时给出nan——分母为 0,需要靠clamp(min=1)或者框架自带的zero_division参数兜底,别指望不会遇到,数据加载器里加了个过滤条件就可能出现空批次。
6.2 MetricCollection 的 key 命名与日志对接
MetricCollection把多个指标打包,update只调一次,内部会把同一份预测分发给每个子指标,省掉了重复的argmax。它的compute()返回一个字典,key 默认是你在构造时给的短名,实际写日志时建议加上prefix和postfix,比如MetricCollection(..., prefix="val/", postfix="_epoch"),避免训练和验证的同名指标在日志系统里互相覆盖。这一点在做对比实验时特别重要,两条曲线画在同一个图上分不清是谁,白做一轮实验。
from torchmetrics import MetricCollection from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score, MulticlassAUROC val_metrics = MetricCollection( { "acc": MulticlassAccuracy(num_classes=10), "f1": MulticlassF1Score(num_classes=10, average="macro"), "auroc": MulticlassAUROC(num_classes=10), }, prefix="val/", ).to(device)注意 AUROC 需要概率、Accuracy 需要 logits 这两个需求冲突时,统一传概率即可——Accuracy 对概率输入同样能处理。
6.3 验证循环里忘记 reset 的代价
这个错误的杀伤力被严重低估了。忘记 reset 的直接后果是验证集指标持续混入历史数据,曲线看起来"异常平滑",越训越好,但模型实际早就过拟合了。更麻烦的是当你的指标状态里存了全部预测(AUROC、AveragePrecision),几个 epoch 不 reset 就足以把显存吃干净,报错信息还是 OOM,你会先去怀疑 batch_size 和模型结构。
我自己的习惯是把reset写在评估函数的最后一行,而不是下一次评估的开头。理由是:写在开头的话,如果某次评估中途异常退出,状态就脏了;写在结尾,函数正常返回时状态一定是干净的,异常时也更容易定位。
最后一个实用技巧是给指标对象加个包装,在compute之后自动reset,把两步合成一步,从源头上消灭忘记调用的可能:
def evaluate(model, loader, metrics, device): model.eval() with torch.no_grad(): for images, targets in loader: images, targets = images.to(device), targets.to(device) probs = torch.softmax(model(images), dim=-1) metrics.update(probs, targets) result = metrics.compute() metrics.reset() return result我个人在实际项目里的体会是,torchmetrics 真正省下的时间不在写指标那几行代码,而在排查"为什么单卡和双卡数字对不上""为什么验证曲线不会掉"这类问题上。它把指标的中间状态显式化了,你想验证对不对,随时可以打印中间状态看累加过程,这比对着一个黑盒函数猜要踏实得多。