如何在PyTorch Lightning中记录每轮中位数损失？解决自定义Metric报错

问题分析 #

你遇到的核心问题：

1. PyTorch Lightning的self.log仅支持min/max/mean/sum作为内置reduce_fx，不直接兼容torch.median；
2. 自定义TorchMetrics子类时，update方法被传入3个参数报错——这是因为TorchMetrics的Metric类默认设计为接收模型预测值（preds）和目标值（targets）两个核心参数，即使你只log单个损失值，Lightning也会自动传入第二个占位参数；
3. 你之前的自定义Metric逻辑有误：每次update直接覆盖中位数，而非累积所有batch损失后计算全局中位数，最终得到的只是最后一个batch的中位数，并非整轮结果。

解决方案 #

方案一：修正自定义TorchMetrics类（推荐，支持分布式）

重新实现MedianMetric，兼容update的参数要求，同时正确累积所有损失再计算全局中位数：

import torch
from torchmetrics import Metric

class MedianMetric(Metric):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # 用dist_reduce_fx="cat"自动处理分布式下的张量聚合
        self.add_state("all_losses", default=[], dist_reduce_fx="cat")
    
    def update(self, loss, *args, **kwargs):
        # 接受额外参数（无需使用），将当前batch的损失加入列表
        self.all_losses.append(loss)
    
    def compute(self):
        # 拼接所有批次损失，计算全局中位数
        all_losses = torch.cat(self.all_losses)
        return torch.median(all_losses).float()
    
    def __name__(self):
        return "median"

在LightningModule中使用：

class MyModel(pl.LightningModule):
    def __init__(self, ...):
        super().__init__()
        self.median_loss_metric = MedianMetric()
        # 其他初始化逻辑...
    
    def training_step(self, batch, batch_idx):
        loss, _ = self.calculate_metrics_and_loss(batch, batch_idx, metric=False)
        # 传入自定义Metric作为reduce_fx，自动完成累积和计算
        self.log(
            "train_loss_median", 
            loss, 
            on_epoch=True, 
            prog_bar=True, 
            reduce_fx=self.median_loss_metric
        )
        # 可同时记录平均损失
        self.log("train_loss_mean", loss, on_epoch=True, prog_bar=True)
        return loss

方案二：手动在training_epoch_end计算（简便，适合非分布式或自定义逻辑）

无需自定义Metric，直接在轮次结束时手动收集所有batch损失并计算中位数：

class MyModel(pl.LightningModule):
    def __init__(self, ...):
        super().__init__()
        self.training_losses = []
        # 其他初始化逻辑...
    
    def training_step(self, batch, batch_idx):
        loss, _ = self.calculate_metrics_and_loss(batch, batch_idx, metric=False)
        # 记录当前batch的损失
        self.training_losses.append(loss)
        self.log("train_loss_mean", loss, on_epoch=True, prog_bar=True)
        return loss
    
    def training_epoch_end(self, outputs):
        # 拼接所有批次损失
        all_losses = torch.cat(self.training_losses)
        
        # 分布式训练需同步所有进程的损失（非分布式可跳过此段）
        if self.trainer.is_distributed:
            gathered_losses = [torch.zeros_like(all_losses) for _ in range(self.trainer.world_size)]
            torch.distributed.all_gather(gathered_losses, all_losses)
            all_losses = torch.cat(gathered_losses)
        
        # 计算并log中位数损失
        median_loss = torch.median(all_losses)
        self.log("train_loss_median", median_loss, prog_bar=True)
        
        # 清空列表，准备下一轮
        self.training_losses.clear()

关键说明 #

• 方案一的自定义Metric会自动处理分布式训练中的张量同步，无需额外代码；
• 方案二更灵活，适合需要在轮次结束时添加其他自定义逻辑的场景；
• 两种方案都不会重复计算损失，均基于训练中已得到的loss值进行统计。

内容的提问来源于stack exchange，提问作者Ryan Dempsey