[FIXED] How to integrate pytorch lightning profiler with tensorboard?

Issue

I know we can use torch profiler with tensorboard using something like this:

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.

It works perfectly with pytorch, but the problem is I have to use pytorch lightning and if I put this in my training step, it just doesn't create the log file nor does it create an entry for profiler. All I get is lightning_logs which isn't the profiler output. I couldn't find anything in the docs about lightning_profiler and tensorboard so does anyone have any idea? Here's what my training function looks like:

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

Solution

You don't have to use raw torch.profiler at all. There is a whole page in Lightning Docs dedicated to Profiling ..

.. and its as easy as passing a trainer flag called profiler like

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

Also, set TensorBoardLogger as your preferred logger as you normally do

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

Answered By - ayandas