Skip to content

Instantly share code, notes, and snippets.

@FindHao
Created July 5, 2022 21:05
Show Gist options
  • Save FindHao/f2f1202ea62bf5ae97a4a94554cc32be to your computer and use it in GitHub Desktop.
Save FindHao/f2f1202ea62bf5ae97a4a94554cc32be to your computer and use it in GitHub Desktop.
pytorch profiler
t0 = time.time_ns()
activity_groups = []
activity_groups.append(profiler.ProfilerActivity.CUDA)
activity_groups.append(profiler.ProfilerActivity.CPU)
profile_detailed=True
with profiler.profile(
schedule=profiler.schedule(wait=0, warmup=0, active=1),
activities=activity_groups,
record_shapes=profile_detailed,
profile_memory=profile_detailed,
with_stack=profile_detailed,
with_flops=profile_detailed,
on_trace_ready=profiler.tensorboard_trace_handler("/tmp/logs/")
) as prof:
# add your code here.
pass
prof.export_chrome_trace("/tmp/a.pt.trace.json")
t1 = time.time_ns()
print('{:<20} {:>20}'.format("Total Wall Time:", "%.3f milliseconds" % ((t1 - t0) / 1_000_000)), sep='')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment