How to profile a model with DLProf on pythorn ligthning?

[Johnsonms]

Hi Guys:
As this docs https://docs.nvidia.com/deeplearning/frameworks/dlprof-user-guide/#profiling_pytorch_pyprof, we need to do the following to capture the traces, but I don't know where I can place "with torch.autograd.profiler.emit_nvtx():" and " nvtx.init(enable_function_stack=True)" in pytorch lightning:

nvtx.init(enable_function_stack=True)
with torch.autograd.profiler.emit_nvtx():
    for epoch in range(1, 2):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

The doc as follow:

When profiling PyTorch models, DLProf uses a python pip package called nvidia_dlprof_pytorch_nvtx to insert the correct NVTX markers. nvidia_dlprof_pytorch_nvtx must first be enabled in the PyTorch Python script before it can work correctly. To enable it, you must add the following lines to your PyTorch network:

import nvidia_dlprof_pytorch_nvtx
nvidia_dlprof_pytorch_nvtx.init()

You should also run the training/inference loop with PyTorch’s NVTX Context Manager with the following:

with torch.autograd.profiler.emit_nvtx():

[akihironitta]

@Johnsonms I've never tried this tool, so there might be extra work required, but simply calling trainer.fit(...) in the context doesn't work?

nvtx.init(enable_function_stack=True)
with torch.autograd.profiler.emit_nvtx():
-    train(args, model, device, train_loader, optimizer, epoch)
-    test(model, device, test_loader)
-    scheduler.step()
+    trainer.fit(model, data)