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