PyTorch Lightning CSVLogger：训练与验证损失为何分行记录？

问题描述 #

使用PyTorch Lightning 2.2.2时，在training_step和validation_step中分别记录训练损失（train_loss）和验证损失（val_loss），生成的metrics.csv会将同一epoch的两个指标分两行记录，示例如下：

同一epoch的步号相同（已省略），想知道能否通过内置CSVLogger将同epoch的两个指标放在同一行。

复现代码 #

import torch
from torch.nn import functional as F
from torch.utils.data import TensorDataset
import lightning as pl
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

iris = load_iris()
features, target = iris.data, iris.target

train_features, val_features, train_target, val_target = train_test_split(features, target, test_size=0.2)

train_features = torch.tensor(train_features).float()
val_features = torch.tensor(val_features).float()
train_target = torch.tensor(train_target).long()
val_target = torch.tensor(val_target).long()

dm = pl.LightningDataModule.from_datasets(
    train_dataset=TensorDataset(train_features, train_target),
    val_dataset=TensorDataset(val_features, val_target),
    batch_size=5,
)

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(4, 3)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = F.cross_entropy(y_hat, y)
        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

    def forward(self, x):
        return self.layer(x)

model = Model()
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, dm)

解决方案 #

可以通过以下两种方式实现：

方法1：在epoch结束时统一记录指标 #

重写training_epoch_end、validation_epoch_end和on_epoch_end方法，先收集epoch级别的平均损失，再在epoch末尾一次性将两个指标一起记录：

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(4, 3)
        self.train_epoch_loss = []
        self.val_epoch_loss = []

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = F.cross_entropy(y_hat, y)
        self.train_epoch_loss.append(loss)
        return loss

    def training_epoch_end(self, outputs):
        avg_train_loss = torch.stack(self.train_epoch_loss).mean()
        self.log("train_loss", avg_train_loss, prog_bar=True)
        self.train_epoch_loss.clear()

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = F.cross_entropy(y_hat, y)
        self.val_epoch_loss.append(loss)

    def validation_epoch_end(self, outputs):
        avg_val_loss = torch.stack(self.val_epoch_loss).mean()
        self.log("val_loss", avg_val_loss, prog_bar=True)
        self.val_epoch_loss.clear()

    def on_epoch_end(self):
        # 合并两个指标并记录，确保同epoch数据在同一行
        self.log_dict(
            {"train_loss": self.trainer.callback_metrics["train_loss"], 
             "val_loss": self.trainer.callback_metrics["val_loss"]},
            sync_dist=True
        )

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

    def forward(self, x):
        return self.layer(x)

方法2：自定义CSVLogger合并逻辑 #

继承官方CSVLogger，重写日志写入逻辑，自动合并同一epoch的指标：

from lightning.pytorch.loggers.csv_logs import CSVLogger, ExperimentWriter

class MergedCSVLogger(CSVLogger):
    def __init__(self, save_dir: str, name: str = "lightning_logs", version: str | int | None = None, prefix: str = "", flush_logs_every_n_steps: int = 100):
        super().__init__(save_dir, name, version, prefix, flush_logs_every_n_steps)

    def _create_experiment_writer(self) -> ExperimentWriter:
        return MergedExperimentWriter(self.log_dir, self._prefix)

class MergedExperimentWriter(ExperimentWriter):
    def __init__(self, log_dir: str, prefix: str = ""):
        super().__init__(log_dir, prefix)
        self.epoch_buffer = {}

    def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        epoch = metrics.get("epoch")
        if epoch is not None:
            # 缓存当前epoch的所有指标
            if epoch not in self.epoch_buffer:
                self.epoch_buffer[epoch] = {}
            self.epoch_buffer[epoch].update(metrics)
            # 当两个指标都收集完成时，再写入文件
            if "train_loss" in self.epoch_buffer[epoch] and "val_loss" in self.epoch_buffer[epoch]:
                merged_metrics = self.epoch_buffer.pop(epoch)
                super().log_metrics(merged_metrics, step)
        else:
            super().log_metrics(metrics, step)

使用自定义Logger启动训练：

model = Model()
trainer = pl.Trainer(max_epochs=10, logger=MergedCSVLogger(save_dir="."))
trainer.fit(model, dm)

说明 #

方法1实现简单，利用PyTorch Lightning的回调机制即可完成；方法2更灵活，适合需要扩展更多指标合并场景的情况。

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