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

Question

我知道我们可以使用带有 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

Answer 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(..))