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_lastFalse会让最后一批的样本数少很多。举个例子前 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_updateTrueforward的内部动作是先把当前累积的状态备份下来reset清空只用这一批数据做一次update加compute拿到结果再把备份的状态恢复回去。所以forward返回的是当前这一批数据单独的指标值跟你之前累积的历史无关。这个语义差异非常关键。训练时想看这个 batch 表现如何用metric(preds, target)验证时想看整个验证集表现如何必须用update累积再compute。两者混用会出现训练日志里的准确率忽高忽低、验证日志里纹丝不动这种诡异现象。另外forward因为有状态备份和恢复的动作开销比update大不少在数据量大、batch 数多的循环里频繁调用会拖慢训练。3. 分类指标选型的判断依据从 Accuracy 到 AUROCtorchmetrics 的分类指标从 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_classes10, averagemacro).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 参数在小类别上的连锁反应averagemicro把所有类别的预测摊平成一个二分类问题算得到的数值通常接近整体准确率看起来漂亮averagemacro对每个类别单独算再取算术平均小类别会被赋予和大类别相同的权重数值往往低一截averageweighted按类别样本数加权介于两者之间。选哪个不是审美问题而是你要回答什么业务问题——想知道模型整体对了多少用 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_fxmean合并在各进程样本数相同的情况下没问题样本数不同时比如最后一个 batch 分配不均或者用了加权采样直接对平均值求平均是错的必须先合并总和与计数最后再除。这就是为什么自定义指标时推荐把中间状态设计成可加的量。4.2 dist_sync_on_step 打开之后发生了什么dist_sync_on_stepTrue会让forward每次调用都触发同步。它解决的是想在训练中途看实时全局指标的需求代价是每个 batch 一次通信。绝大多数情况下不需要打开它——你真要实时看趋势也应该在固定间隔才记录一次而不是每个 step 都同步。还有一种情况是同步了但结果还是不对指标对象没有.to(device)。如果模型在 GPU 上、指标对象还留在 CPUupdate时框架会把输入张量搬到指标所在的设备这个隐式拷贝不会报错但会悄悄拖慢速度并且当你的状态张量很大比如保存全部预测用于算 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_cpuTrue之后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, weightsNone, **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, defaulttorch.zeros(num_classes), dist_reduce_fxsum) self.add_state(pos, defaulttorch.zeros(num_classes), dist_reduce_fxsum) 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(min1) w self.weights.to(recall.device) return (recall * w).sum() / w.sum()这段代码把加权召回率拆成了两个可加的计数器dist_reduce_fxsum保证多卡合并时分子分母分别相加比值自然正确。如果图省事直接把每个 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, dtypetorch.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(min1)或者框架自带的zero_division参数兜底别指望不会遇到数据加载器里加了个过滤条件就可能出现空批次。6.2 MetricCollection 的 key 命名与日志对接MetricCollection把多个指标打包update只调一次内部会把同一份预测分发给每个子指标省掉了重复的argmax。它的compute()返回一个字典key 默认是你在构造时给的短名实际写日志时建议加上prefix和postfix比如MetricCollection(..., prefixval/, postfix_epoch)避免训练和验证的同名指标在日志系统里互相覆盖。这一点在做对比实验时特别重要两条曲线画在同一个图上分不清是谁白做一轮实验。from torchmetrics import MetricCollection from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score, MulticlassAUROC val_metrics MetricCollection( { acc: MulticlassAccuracy(num_classes10), f1: MulticlassF1Score(num_classes10, averagemacro), auroc: MulticlassAUROC(num_classes10), }, prefixval/, ).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 真正省下的时间不在写指标那几行代码而在排查为什么单卡和双卡数字对不上为什么验证曲线不会掉这类问题上。它把指标的中间状态显式化了你想验证对不对随时可以打印中间状态看累加过程这比对着一个黑盒函数猜要踏实得多。