PyTorch Lightning 自定义 Profiler 完全指南构建自己的性能分析器与剖析自定义代码段【免费下载链接】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 性能分析Profiling主题的进阶篇面向希望构建自定义 Profiler或对训练流程中特定代码片段进行精细剖析的开发者。通过继承lightning.pytorch.profilers.profiler.Profiler抽象基类并重写start、stop、summary等方法你可以完全掌控剖析的记录逻辑与报告格式同时借助self.profiler.profile()上下文管理器可在 LightningModule 的任何位置为自定义动作计时。读完本文你将掌握自定义 Profiler 的完整编写与接入流程、profile()上下文管理器的底层原理以及内置五种 Profiler 的选型依据。Profiler 体系速览五种内置实现与统一基类在动手自定义之前先明确 Lightning 的 Profiler 家族结构。所有 Profiler 都继承自Profiler抽象基类见 src/lightning/pytorch/profilers/profiler.py统一从 src/lightning/pytorch/profilers/init.py 导出Profiler底层机制适用场景PassThroughProfiler空实现start/stop直接passTrainer 默认使用零开销SimpleProfilertime.perf_counter()计时记录各动作平均耗时与总耗时AdvancedProfilerPython 标准库cProfile逐函数调用级明细输出冗长PyTorchProfilertorch.profiler/ autograd Profiler算子级 CPU/GPU 开销分析XLAProfilerXLA 设备事件TPU 场景其中PassThroughProfiler的源码src/lightning/pytorch/profilers/base.py印证了其设计意图当你不需要哪怕很小的剖析开销时使用此类Trainer 默认使用它。这意味着自定义 Profiler 的价值在于针对你的训练瓶颈给出内置方案之外的定制观测维度。构建你自己的 Profiler继承抽象基类构建自定义 Profiler 的核心方法是继承Profiler并重写其抽象方法。Profiler基类定义了两个必须实现的抽象方法src/lightning/pytorch/profilers/profiler.pystart(action_name)定义如何开始记录一个动作stop(action_name)定义动作完成时如何记录耗时。此外还有三个可选的钩子方法summary()默认返回空字符串重写后生成报告文本setup(stage, local_rank, log_dir)训练前的初始化钩子负责注入阶段名、local rank 与日志目录src/lightning/pytorch/profilers/profiler.pyteardown(stage)训练结束后的清理钩子默认关闭已打开的输出文件与写流src/lightning/pytorch/profilers/profiler.py。完整示例统计动作出现次数与首次发生时间的 ActionCountProfiler官方文档给出了一个统计每个动作首次发生时间与总调用次数的自定义 Profiler。下面是在保留原示例逻辑的基础上修正了原文档 f-string 中缺失的格式化前缀使其可直接运行from lightning.pytorch.profilers import Profiler from collections import defaultdict import time class ActionCountProfiler(Profiler): def __init__(self, dirpathNone, filenameNone): super().__init__(dirpathdirpath, filenamefilename) self._action_count defaultdict(int) self._action_first_occurrence {} def start(self, action_name): if action_name not in self._action_first_occurrence: self._action_first_occurrence[action_name] time.strftime(%m/%d/%Y, %H:%M:%S) def stop(self, action_name): self._action_count[action_name] 1 def summary(self): res f\nProfile Summary: \n max_len max(len(x) for x in self._action_count) for action_name in self._action_count: # generate summary for actions called more than once if self._action_count[action_name] 1: res ( f{action_name:{max_len}s} \t f{self._action_first_occurrence[action_name]} \t f{self._action_count[action_name]} \n ) return res def teardown(self, stage): self._action_count {} self._action_first_occurrence {} super().teardown(stagestage)接入方式与内置 Profiler 完全一致——把实例传给 Trainer 的profiler参数即可trainer Trainer(profilerActionCountProfiler()) trainer.fit(...)实现要点剖析对照基类源码这个示例触及了几个关键设计点1. 构造参数约定。基类__init__接收dirpath报告输出目录与filename报告文件名自动补.txt后缀并维护_output_file、_write_stream、_local_rank、_stage等内部状态src/lightning/pytorch/profilers/profiler.py。自定义子类应通过super().__init__(dirpathdirpath, filenamefilename)保留这一约定让报告输出能力开箱即用。2. 输出流的两条路径。基类_prepare_streams()src/lightning/pytorch/profilers/profiler.py决定报告去向当同时提供了filename与dirpath时写入文件使用lightning.fabric.utilities.cloud_io.get_filesystem以兼容云存储否则通过_rank_zero_info打印到日志且只在 rank 0或未设置 rank时输出避免多进程下重复打印。3.describe()统一出口。训练结束后由 Trainer 调用describe()src/lightning/pytorch/profilers/profiler.py其流程是_prepare_streams()→ 取summary()结果写入流 →flush()→teardown()。因此你的自定义summary()只需返回字符串文件落盘与打印都由基类托管。4.teardown的双重职责。示例在teardown中清空统计字典并调用super().teardown(stagestage)关闭文件句柄——这一步不可省略否则会泄漏输出文件句柄。基类的__del__也会兜底调用teardownsrc/lightning/pytorch/profilers/profiler.py。剖析感兴趣的自定义动作profile() 上下文管理器除了剖析训练循环中 Lightning 自动埋点的动作如training_step、validation_step你还可以在 LightningModule 内引用 profiler对任意自定义代码段计时。其标准写法分两步。第一步在模块中保存 profiler 引用from lightning.pytorch.profilers import SimpleProfiler, PassThroughProfiler class MyModel(LightningModule): def __init__(self, profilerNone): self.profiler profiler or PassThroughProfiler()这里用PassThroughProfiler()作为默认值保证不传 profiler 时零开销——这正是PassThroughProfiler的设计定位。若传入SimpleProfiler()则自定义动作的耗时会被真实记录。第二步用with self.profiler.profile(...)包裹代码段class MyModel(LightningModule): def custom_processing_step(self, data): with self.profiler.profile(my_custom_action): ... return dataprofile()是基类提供的上下文管理器src/lightning/pytorch/profilers/profiler.py其实现保证了生命周期安全contextmanager def profile(self, action_name: str) - Generator: try: self.start(action_name) yield action_name finally: self.stop(action_name)进入with块即调用start(action_name)无论块内是否抛出异常退出时都会通过finally确保调用stop(action_name)——因此即使剖析的代码段崩溃也不会造成计时器泄漏或只 start 不 stop的状态残留。你可以把这段with self.profiler.profile(action_name)用在模块的前向、数据处理、损失计算等任何位置。完整可运行代码将 profiler 同时传入模块与 Trainer同一个实例即可让剖析动作贯穿整个训练流程from lightning.pytorch.profilers import SimpleProfiler, PassThroughProfiler class MyModel(LightningModule): def __init__(self, profilerNone): self.profiler profiler or PassThroughProfiler() def custom_processing_step(self, data): with self.profiler.profile(my_custom_action): ... return data profiler SimpleProfiler() model MyModel(profiler) trainer Trainer(profilerprofiler, max_epochs1)注意MyModel(profiler)与Trainer(profilerprofiler)必须共享同一实例Trainer 会通过它内部的连接器把 profiler 注入 LightningModule从源码看PyTorchProfiler内部通过_lightning_module属性持有模块引用见 src/lightning/pytorch/profilers/pytorch.py自定义 profiler 应保持同样的单一实例原则否则模块内计时与训练循环计时会落在两套独立状态上。深入理解报告生成SimpleProfiler 的参考实现若你的自定义 Profiler 需要输出结构化报告SimpleProfiler是最佳参考模板src/lightning/pytorch/profilers/simple.py。它演示了 Profiler 子类的完整工程实践1. 计时数据结构。current_actions: dict[str, float]记录正在计时的动作及其time.perf_counter()起始时刻recorded_durations: dict按动作名累积每次耗时。2. 健壮性校验。start时若动作已在计时中则抛ValueErrorstop时若动作从未启动同样抛ValueErrorsrc/lightning/pytorch/profilers/simple.py。这能尽早暴露嵌套同名with块等误用。3. 可扩展报告。构造函数提供extendedTrue开关src/lightning/pytorch/profilers/simple.py扩展模式下报告每列的 Mean duration (s) / Num calls / Total time (s) / Percentage %并按耗时占比降序排列关闭后仅保留均值与总计两列。这种核心逻辑 可选列的结构同样适用于自定义报告。4. 浮点精度。统计时使用math.fsumsrc/lightning/pytorch/profilers/simple.py累加大量微小时耗避免普通求和累积误差。选型建议与自定义场景判断结合五种内置 Profiler 与自定义能力可按下述思路决策只是快速定位训练循环瓶颈直接用SimpleProfiler其extendedTrue报告会标出耗时占比最高的动作需要算子级或 GPU 时间使用PyTorchProfiler其sort_by_key支持cuda_time_total、cuda_memory_usage等键完整合法键集见 src/lightning/pytorch/profilers/pytorch.py需要逐函数调用明细使用AdvancedProfiler基于cProfile注意其输出冗长且dump_statsTrue时必须提供dirpathsrc/lightning/pytorch/profilers/advanced.py内置方案无法覆盖你的观测维度例如统计动作调用次数、首次发生时间、自定义指标采样继承Profiler自定义并复用基类的profile()、describe()、文件输出能力模块内的局部代码段计时不必自定义 Profiler直接用self.profiler.profile(action)配合任意内置 Profiler 即可。验证与测试参考现有测试用例仓库的 profiler 测试tests/tests_pytorch/profilers/test_profiler.py是验证自定义实现是否合规的现成参照SimpleProfiler的start/stop配对与异常路径均有用例覆盖如 test_profiler.pywith profiler.profile(action)上下文管理器在训练钩子on_test_start等与嵌套场景下均被测试如 test_profiler.py 处的嵌套profile用例多 Profiler 的构建参数序列化__reduce__也有对应覆盖test_profiler.py自定义 Profiler 若含不可 pickle 的成员如cProfile.Profile可参考AdvancedProfiler.__reduce__src/lightning/pytorch/profilers/advanced.py提供安全的序列化路径保证分布式场景下能正确复制。结语自定义 Profiler 的本质是接入 Lightning 统一的动作计时 报告输出框架继承Profiler、重写start/stop/summary、交给 Trainer 调度即可在零侵入的前提下获得完全定制化的性能观测而self.profiler.profile()则把这一能力下沉到模型代码的任意角落。将本文示例与仓库中的SimpleProfiler、AdvancedProfiler实现对照阅读你便能快速设计出贴合自身训练瓶颈分析需求的剖析工具。【免费下载链接】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），仅供参考
