PyTorch Lightning进度条设置：实现每个Epoch换行显示

直接设置TQDM的leave=True无效，是因为PyTorch Lightning的默认进度条回调会在每个Epoch结束后重置或关闭进度条，覆盖了你的参数。可以通过以下两种方法解决：

方法一：自定义进度条回调 #

继承PyTorch Lightning的ProgressBar类，重写初始化TQDM的方法，强制设置leave=True：

from pytorch_lightning.callbacks import ProgressBar
from tqdm import tqdm

class NewLineProgressBar(ProgressBar):
    def init_train_tqdm(self):
        # 初始化训练进度条，设置leave=True保留进度条
        return tqdm(
            desc="Training",
            total=self.max_steps,
            leave=True,
            dynamic_ncols=True,
            file=self.train_progress_bar.file,
            disable=self.is_disabled,
        )

    def init_validation_tqdm(self):
        # 初始化验证进度条，设置leave=True保留进度条
        return tqdm(
            desc="Validating",
            leave=True,
            dynamic_ncols=True,
            file=self.val_progress_bar.file,
            disable=self.is_disabled,
        )

然后在Trainer中使用这个自定义回调：

trainer = Trainer(callbacks=[NewLineProgressBar()])

方法二：禁用默认进度条，手动控制 #

如果自定义回调太麻烦，可以直接禁用Lightning的默认进度条，自己手动添加TQDM进度条，完全控制显示逻辑：

from tqdm import tqdm
import pytorch_lightning as pl

# 禁用默认进度条
trainer = pl.Trainer(enable_progress_bar=False, max_epochs=10)

# 获取数据加载器
train_loader = trainer.train_dataloader
val_loader = trainer.val_dataloaders[0]

for epoch in range(trainer.max_epochs):
    print(f"\nEpoch {epoch+1}/{trainer.max_epochs}")
    
    # 训练循环进度条
    train_bar = tqdm(train_loader, leave=True, desc="Training")
    for batch_idx, batch in enumerate(train_bar):
        trainer.model.train()
        loss = trainer.model.training_step(batch, batch_idx)
        # 更新进度条后缀显示loss
        train_bar.set_postfix({"Train Loss": f"{loss.item():.4f}"})
    
    # 验证循环进度条
    val_bar = tqdm(val_loader, leave=True, desc="Validating")
    for batch_idx, batch in enumerate(val_bar):
        trainer.model.eval()
        val_loss = trainer.model.validation_step(batch, batch_idx)
        val_bar.set_postfix({"Val Loss": f"{val_loss.item():.4f}"})

这两种方法都能让每个Epoch的进度条保留在新行，不会被覆盖。第一种方法更贴合Lightning的原生逻辑，第二种方法自由度更高。

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