
PyTorch Lightning 进度条定制完全指南TQDM 与 Rich 的实现原理与实战配置【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning导读本文将系统讲解 PyTorch Lightning 中进度条Progress Bar的定制方法。Lightning 内置了基于tqdm的TQDMProgressBar默认使用与基于rich的RichProgressBar两套实现并提供了可自由扩展的基类ProgressBar。读完本文你将掌握如何在Trainer中切换进度条风格、调整刷新频率与显示时机、自定义进度条样式与文本列、从零实现自己的进度条回调并理解进度条在源码层面是如何被自动装配与驱动的。进度条体系概览三套层次分明的 APIPyTorch Lightning 的进度条实现分为三个层级全部位于 src/lightning/pytorch/callbacks/progress/ 目录下类说明文件ProgressBar所有进度条的基类继承自Callback内置批次总数计算与多 DataLoader 追踪逻辑progress_bar.pyTQDMProgressBar基于tqdm的默认进度条打印到stdouttqdm_progress.pyRichProgressBar基于rich的富文本进度条支持主题与自定义列rich_progress.py三者通过Trainer的callbacks参数统一接入。由于TQDMProgressBar是默认进度条你甚至无需显式传递任何回调参数即可获得训练进度显示。Trainer 如何自动装配进度条从源码看进度条并非由用户手工挂在 Trainer 上而是由_CallbackConnector在Trainer初始化阶段自动完成的核心逻辑位于 callback_connector.py 的_configure_progress_bar如果callbacks列表中已经包含多个ProgressBar子类实例会抛出MisconfigurationException——同一个 Trainer 同时只支持一个进度条如果用户显式传入了进度条回调且enable_progress_barTrue则直接使用用户提供的实例不再添加默认进度条如果用户传入了进度条回调但设置了enable_progress_barFalse同样会抛出MisconfigurationException提示配置冲突如果两者都未指定则当环境中安装了rich时默认使用RichProgressBar否则回退到TQDMProgressBar见 callback_connector.py。Trainer还暴露了progress_bar_callback只读属性用于在运行期取回当前生效的进度条实例见 trainer.py。进度条在分布式训练下的自动降级基类ProgressBar.setup中有个重要细节当trainer.is_global_zero为假即非 rank 0 进程时会自动调用disable()关闭进度条显示避免多进程下终端输出互相覆盖见 progress_bar.py。TQDMProgressBar默认进度条的进阶使用TQDMProgressBar内部使用tqdm库打印到stdout最多同时展示四条进度条sanity check 进度验证集 sanity check 运行期间的进度训练进度显示训练进度验证开始时暂停、结束时恢复并能在使用Trainer(val_check_interval...)时正确应对训练中多次穿插验证的情况验证进度仅在验证期间可见展示所有验证 DataLoader 的总体进度测试进度仅在测试时激活展示所有测试 DataLoader 的总体进度。此外源码中还实现了predict阶段的进度条init_predict_tqdm供trainer.predict()使用。对于无限大小的数据集如IterableDataset进度条永远不会到达 100%源码通过convert_inf将inf转为None见 tqdm_progress.py。构造参数与默认值从 tqdm_progress.py 的构造函数可以看到三个公开参数参数默认值含义refresh_rate1每处理多少个 batch 刷新一次进度条设为0可完全禁用显示process_position0进度条的起始偏移行数0 时用于与脚本中其他 tqdm 进度条共存leaveFalse每个 epoch 结束后是否在终端保留该条进度条True表示保留调整刷新频率from lightning.pytorch.callbacks import TQDMProgressBar trainer Trainer(callbacks[TQDMProgressBar(refresh_rate10)])refresh_rate的语义是“每隔 N 个 batch 刷新一次”其判定逻辑见 tqdm_progress.pydef _should_update(self, current: int, total: int) - bool: return self.is_enabled and (current % self.refresh_rate 0 or current total)即当current是refresh_rate的整数倍或者到达最后一个 batchcurrent total时才会触发刷新。配套的测试 test_tqdm_progress_bar.py 中的test_tqdm_progress_bar_progress_refresh与test_train_progress_bar_update_amount验证了不同refresh_rate下刷新次数的精确行为。值得一提的是refresh_rate还支持两个环境变量层面的特殊处理见_resolve_refresh_ratetqdm_progress.py在 Google Colab 环境检测到COLAB_GPU环境变量且refresh_rate保持默认值1时会自动提升为20以避免高频刷新导致崩溃若设置了TQDM_MINITERS环境变量则实际刷新率取max(int(TQDM_MINITERS), refresh_rate)即允许外部强制提高最小刷新间隔。每个 epoch 结束时保留进度条默认情况下训练进度条在每个新 epoch 开始时会被重置覆盖。若希望在每个 epoch 结束时在终端保留一条历史进度条将leave置为Truetrainer Trainer(callbacks[TQDMProgressBar(leaveTrue)])源码中on_train_epoch_starttqdm_progress.py会在_leave为真时重新init_train_tqdm()创建新进度条而on_train_epoch_end则负责关闭旧进度条正是这一对钩子共同实现了“每个 epoch 保留一条独立进度条”的效果。通过覆写方法定制默认 tqdm 行为TQDMProgressBar将每条进度条的创建逻辑收敛到独立的init_*_tqdm方法中覆写即可定制描述文本等外观from lightning.pytorch.callbacks import TQDMProgressBar class LitProgressBar(TQDMProgressBar): def init_validation_tqdm(self): bar super().init_validation_tqdm() bar.set_description(running validation...) return bar trainer Trainer(callbacks[LitProgressBar()])可覆写的工厂方法包括init_sanity_tqdm、init_train_tqdm、init_validation_tqdm、init_test_tqdm、init_predict_tqdm。每条进度条都统一使用类属性BAR_FORMAT定义格式串tqdm_progress.py。关于smoothing参数的注意事项官方文档特别强调在使用默认的TQDMProgressBar实现时smoothing选项不生效。原因在于 Lightning 更新进度条使用的是bar.refresh()而非bar.update()——从源码的_update_n可见tqdm_progress.py它直接设置bar.n value再调用bar.refresh()。这种方式可能使进度条与实际进度产生细微的不同步。如果需要用到smoothing可以覆写TQDMProgressBar改用bar.update()方法驱动。此外项目自定义的Tqdm类重写了format_num会对浮点数字符串进行右端补零填充_PAD_SIZE 5以消除进度条因数字宽度抖动而产生的闪烁对应测试见 test_tqdm_progress_bar.py 中的test_tqdm_format_num。RichProgressBar更美观的富文本进度条Rich是一个用于终端富文本与精美格式化的 Python 库。RichProgressBar相比TQDMProgressBar提供了主题配色、自定义列等更灵活的外观控制能力。安装与启用首先安装rich依赖pip install rich然后在Trainer中启用from lightning.pytorch.callbacks import RichProgressBar trainer Trainer(callbacks[RichProgressBar()])注意如果环境未安装rich且版本低于10.2.2RichProgressBar的构造函数会直接抛出ModuleNotFoundError见 rich_progress.py。也正因为此Trainer 在检测到rich可用时才会将其选为默认进度条。构造参数从 rich_progress.py 可以看到RichProgressBar的四个参数注意其refresh_rate语义与 tqdm 版本不同参数默认值含义refresh_rate100每秒刷新的次数频率而不是按 batch 数设为0禁用显示leaveFalse每个 epoch 结束时是否在终端保留进度条themeRichProgressBarTheme()控制各组件配色的主题对象console_kwargsNone透传给rich.console.Console的参数字典自定义主题配色RichProgressBarTheme是一个dataclass其字段与默认值如下rich_progress.py字段默认值作用description描述文字样式如 Epoch x、Testingprogress_bar#6206E0进度条填充部分的样式progress_bar_finished#6206E0进度条完成后的样式progress_bar_pulse#6206E0处理IterableDataset无限大小时的脉冲动画样式batch_progress批次进度文本如 10/50样式timedim已用时间与剩余时间文本样式processing_speeddim underline处理速度it/s文本样式metricsitalic指标文本样式metrics_text_delimiter 多个指标之间的分隔符metrics_format.3f指标的数值格式化串所有颜色值遵循 Rich 的样式语法十六进制色值或命名颜色。一个完整的自定义示例from lightning.pytorch.callbacks import RichProgressBar from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBarTheme # create your own theme! progress_bar RichProgressBar( themeRichProgressBarTheme( descriptiongreen_yellow, progress_bargreen1, progress_bar_finishedgreen1, progress_bar_pulse#6206E0, batch_progressgreen_yellow, timegrey82, processing_speedgrey82, metricsgrey82, metrics_text_delimiter\n, metrics_format.3e, ) ) trainer Trainer(callbacksprogress_bar)这里的metrics_format控制指标数值的展示精度例如.3e表示科学计数法保留 3 位小数默认的.3f表示保留 3 位小数。配套测试 test_rich_progress_bar.py 中的test_rich_progressbegin▁of▁sentenceimport_metrics_format与test_rich_progress_bar_metrics_theme_update验证了这些字段被正确传递到 Rich 组件中。自定义进度条列configure_columnsRichProgressBar将进度条拆分为多个“列”column默认列由configure_columns方法返回rich_progress.py依次为TextColumn描述文本如 Epoch 0CustomBarColumn进度条本体支持主题配色与无限数据集脉冲动画BatchesProcessedColumn批次进度计数如 10/50CustomTimeColumn已用时间与剩余时间如 00:0301:37ProcessingSpeedColumn处理速度如 10.84it/s。覆写configure_columns可以整体替换这些列from rich.progress import TextColumn from lightning.pytorch.callbacks import RichProgressBar custom_column TextColumn([progress.description]Custom Rich Progress Bar!) class CustomRichProgressBar(RichProgressBar): def configure_columns(self, trainer): return [custom_column] progress_bar CustomRichProgressBar()注意覆写configure_columns后默认指标列MetricsTextColumn仍会被_init_progress自动追加到列列表末尾rich_progress.py因此替换列不影响指标展示。测试 test_rich_progress_bar.py 中的test_rich_progress_bar_configure_columns验证了自定义列机制。每 epoch 保留进度条与 tqdm 版本一致RichProgressBar同样通过leaveTrue在每个 epoch 结束时保留一条历史进度条from lightning.pytorch.callbacks import RichProgressBar trainer Trainer(callbacks[RichProgressBar(leaveTrue)])从源码看on_train_epoch_start在_leave为真时会先_stop_progress()停止当前进度再重新_init_progress(trainer)开启新进度rich_progress.py。测试test_rich_progress_bar_leave通过参数化(leave, reset_call_count)精确验证了这一行为。对无限数据集与异常场景的底层支持Rich 实现针对IterableDataset无长度、总 batch 数为inf做了专门适配CustomBarColumn在remaining非有限值时以脉冲动画pulse代替固定进度条rich_progress.pyCustomProgress.add_task对无限 total 创建CustomInfiniteTask其time_remaining恒为None。同时on_exception与teardown钩子会在训练异常或结束时主动停止进度刷新线程避免终端残留rich_progress.py。测试 test_rich_progress_bar.py 中的test_rich_progress_bar_infinite_epochs、test_rich_progress_bar_keyboard_interrupt分别覆盖了这两类场景。另外Rich 官方提示PyCharm 用户需要在其 Run/Debug 配置中开启 Emulate terminal 选项才能正常显示富文本样式输出。用 ProgressBar 基类实现完全自定义的进度条如果内置的两种进度条都无法满足需求可以直接继承ProgressBar基类。该基类本身是一个Callback并预置了丰富的“基础设施”批次总数属性total_train_batches、total_val_batches_current_dataloader、total_test_batches_current_dataloader、total_predict_batches_current_dataloader这些属性会随 epoch、sanity check、limit_*_batches等动态变化且对无限数据集返回inf见 progress_bar.py多 DataLoader 追踪has_dataloader_changed与reset_dataloader_idx_tracker帮助判断当前批次属于第几个 DataLoader从而正确处理多验证集/多测试集的总数切换内置描述文本属性sanity_check_description、train_description、validation_description、test_description、predict_description指标合并get_metrics将 Trainer 的标准指标如日志版本号v_num与通过self.log(..., prog_barTrue)记录的指标合并输出出现重名时会通过rank_zero_warn发出警告progress_bar.py钩子约定disable()、enable()、print()需要子类实现——disable用于关闭进度条Trainer 在非 rank 0 进程及学习率查找等场景会调用print用于在不破坏进度条排版的前提下打印文本。基类文档中给出了一个最简单的自定义示例progress_bar.pyclass LitProgressBar(ProgressBar): def __init__(self): super().__init__() # dont forget this :) self.enable True def disable(self): self.enable False def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) # dont forget this :) percent (batch_idx / self.total_train_batches) * 100 sys.stdout.flush() sys.stdout.write(f{percent:.01f} percent complete \r) bar LitProgressBar() trainer Trainer(callbacks[bar])实现要点有二构造函数中必须调用super().__init__()覆写训练钩子时保留对父类方法的调用以维持 Trainer 内部状态的一致性。常见操作速查完全禁用进度条进度条默认由 Trainer 自动启用。要禁用直接设置 Trainer 参数trainer Trainer(enable_progress_barFalse)注意该参数在Trainer.__init__中默认值为None实际会在初始化时被解析为True见 trainer.py。若同时又在callbacks中显式传入了进度条回调Trainer 会抛出MisconfigurationException提示冲突callback_connector.py。在不破坏进度条的前提下打印信息训练循环中直接print()可能打断 tqdm/rich 的渲染。推荐使用进度条回调自带的print方法trainer.progress_bar_callback.print(your message)TQDMProgressBar的print会找到当前激活的进度条并调用其write方法tqdm_progress.pyRich 版本则依托 Console 渲染体系二者都能保证输出不破坏进度条。将指标显示在进度条上在LightningModule中调用self.log(metric_name, value, prog_barTrue)即可将指定指标实时展示在进度条尾部多个指标会以metrics_text_delimiterRich 默认空格tqdm 默认逗号分隔拼接。指标渲染前Rich 版本会通过apply_to_collection将torch.Tensor统一转换为 Python 浮点数rich_progress.py避免终端输出张量对象。结语与延伸阅读进度条是 Lightning 训练闭环中最直观的“仪表盘”。通过TQDMProgressBar、RichProgressBar与ProgressBar三层 API你既可以在不写一行定制代码的情况下完成风格切换也可以深入到刷新策略、列组件、主题配色层面打造完全贴合自己工作流的进度显示。感兴趣的读者可以进一步阅读仓库中的相关资源进度条回调源码progress/ 目录Trainer 侧装配逻辑callback_connector.py行为验证测试test_tqdm_progress_bar.py 与 test_rich_progress_bar.py与之配合的模型摘要表定制RichModelSummary富文本环境下 Trainer 默认使用RichModelSummary展示模型结构与进度条同源的self.log(..., prog_barTrue)日志机制logging 指南【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考