如何将 pytorch 闪电分析器与 tensorboard 集成？

Posted 2023-03-27

【中文标题】如何将 pytorch 闪电分析器与 tensorboard 集成？

【英文标题】：How to integrate pytorch lightning profiler with tensorboard?

【发布时间】：2021-10-30 23:38:49

【问题描述】：

我知道我们可以使用带有 tensorboard 的 Torch 分析器，如下所示：

with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        with_stack=True
) as prof:
    for step, batch_data in enumerate(train_loader):
        if step >= (1 + 1 + 3) * 2:
            break
        train(batch_data)
        prof.step()  # Need to call this at the end of each step to notify profiler of steps' boundary.

它与 pytorch 完美配合，但问题是我必须使用 pytorch 闪电，如果我把它放在我的训练步骤中，它不会创建日志文件，也不会为探查器创建条目。我得到的只是lightning_logs，这不是分析器输出。我在文档中找不到有关 Lightning_profiler 和 tensorboard 的任何内容，所以有人知道吗？ 这是我的训练函数的样子：

def training_step(self, train_batch, batch_idx):
        
        with torch.profiler.profile(
        activities=[ProfilerActivity.CPU],
        schedule=torch.profiler.schedule(
            wait=1,
            warmup=1,
            active=2,
            repeat=1),
        with_stack=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
        ) as profiler:
        
            x, y = train_batch
            x = x.float()
          
            logits = self.forward(x) 
            
            loss = self.loss_fn(logits, y)
            profiler.step()
        return loss

【参考方案1】：

您根本不必使用原始的torch.profiler。在 Lightning Docs 中有一个 whole page 专门用于分析 ..

.. 就像 passing a trainer flag 调用 profiler 一样简单

# other profilers are "simple", "advanced" etc
trainer = pl.Trainer(profiler="pytorch")

另外，像往常一样将TensorBoardLogger 设置为首选记录器

trainer = pl.Trainer(profiler="pytorch", logger=TensorBoardLogger(..))

【讨论】：

但我想使用 torch profiler 以便跟踪正在执行的每个操作。我尝试使用trainer = pl.Trainer(profiler="pytorch", logger=TensorBoardLogger(..))，但分析器没有显示在张量板上

@Madara 我猜你还没有将torch_tb_profiler 安装为mentioned here

它工作了吗@马达拉？

是的，它确实做到了@ayandas :)