Proper way to enable and disable autograd profiler

#1

Currently I use the following. Is there a better way to enable it without manually calling __enter__? Is it necessary (I came up with it when it seemed necessary, but now it was maybe refactored?)?

if args.profile_autograd:
    autograd_profiler = torch.autograd.profiler.profile()
    autograd_profiler.__enter__()

# model running

if args.profile_autograd:
    autograd_profiler.__exit__(None, None, None)
    autograd_profiler.export_chrome_trace(args.profile_autograd)

I don’t want to use with construct because I want to keep enabling the profiler under the flag and prefer not to factor out the model code in a separate function. enable()-kind of API exists for autograd itself, so I thought maybe it exists for the profiler as well. It also exists for nvprof: torch.cuda.profiler.start()

#2

Doesn’t profile take an enabled optional argument?

with torch.autograd.profiler.profile(enabled=args.profile_autograd):

1 Like
#3

It would start then without with construct also, right?

autograd_profiler = torch.autograd.profiler.profile(enabled=args.profile_autograd)
# model code
autograd_profiler.export_chrome_trace(args.profile_autograd)

Do you know if I can stop it explicitly? I could not find a method for that in docs: https://pytorch.org/docs/stable/autograd.html#torch.autograd.profiler.profile

#4

No, you still want the context manager.