TensorBoard加载PyTorch Profiler分析结果时持续加载挂住无响应

开发者阿泽

2026-8-29

问题背景

使用torch.profiler.profile采集PyTorch训练性能数据，配合TensorBoard插件查看结果时页面持续加载超过1小时无法展示，采集生成的trace文件单个体积接近10GB。
使用的采集代码如下：

with SummaryWriter(tb_dir) as writer, open(logfn, "wt", encoding="utf-8") as logfp, \
    torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        schedule=torch.profiler.schedule(
            wait=1,
            warmup=1,
            active=2
        ),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(profile_dir, worker_name="hmm")
    ) as profiler:
    while (args.max_epochs is None or epoch < args.max_epochs) and (args.max_time is None or total_elapsed < args.max_time):
        train(model, epoch, train_dataset, ema_node_loss, opt, sched, crit) # 外部定义的训练循环函数
        # 省略其余无关逻辑
        profiler.step()

生成的trace文件信息：

-rw-rw-r-- 1 binesh binesh 10049360679 Jun 20 19:40 hmm.1655768373559.pt.trace.json
-rw-rw-r-- 1 binesh binesh 10094463514 Jun 20 20:54 hmm.1655772846630.pt.trace.json

根因分析

profiler.step()调用位置完全错误。profiler的schedule规则（wait/warmup/active的计数）是基于profiler.step()的调用次数生效的，设计上这个方法需要每跑完一个单batch训练迭代就调用一次。当前代码把profiler.step()放在了epoch层级的循环里，等于每跑完一整个epoch（成百上千个batch）才记一次step，配置的wait=1,warmup=1,active=2逻辑完全失效，profiler全量采集了所有epoch、所有batch的CPU、CUDA事件，直接把trace文件撑到了10G级别，远超浏览器的内存加载上限，必然出现永久加载的问题。
未配置冗余信息裁剪开关，默认开启的shape记录、调用栈采集、内存统计等能力会进一步放大trace体积，加剧加载负担。

修复方案

调整profiler.step()的调用位置：将该方法移动到训练逻辑内部的单batch循环中，保证每完成一个batch的前向、反向、参数更新流程就调用一次，让schedule规则正常生效——仅跳过1个batch、预热1个batch、采集2个batch的性能数据，正常情况下生成的trace文件体积会在几十到几百MB区间，完全在TensorBoard插件的加载能力范围内。
修正后的训练逻辑参考：

def train(model, epoch, dataloader, opt, crit, profiler):
    model.train()
    for batch_data, batch_label in dataloader:
        opt.zero_grad()
        batch_data = batch_data.cuda(non_blocking=True)
        batch_label = batch_label.cuda(non_blocking=True)
        pred = model(batch_data)
        loss = crit(pred, batch_label)
        loss.backward()
        opt.step()
        # 单batch训练完成后调用，对齐profiler的schedule计数
        profiler.step()

按需关闭非必要的采集项压缩体积：初始化torch.profiler.profile时，不需要记录算子输入shape就加record_shapes=False，不需要采集Python调用栈就加with_stack=False，不需要做细粒度显存占用分析就加profile_memory=False，可以在不影响核心性能分析结果的前提下把trace体积再压缩50%以上。
清理现有无效的大体积trace：启动TensorBoard前先把profile_dir目录下生成的10G级json文件移出目录，避免插件扫描到异常大文件持续尝试解析加载。
如果有长周期性能统计需求，不需要依赖TensorBoard加载trace，可以直接在训练逻辑中调用profiler.key_averages().table(sort_by="cuda_time_total", row_limit=20)，直接在训练日志中输出CUDA/CPU耗时最高的算子排行，不需要加载全量trace也能拿到核心性能瓶颈数据。

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