with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule(wait=0, warmup=0, active=6, repeat=1), record_shapes=True, profile_memory=True, with_stack=True, on_trace_ready=trace_handler, ) as prof: # Run the PyTorch Model inside the profile context. for _ inrange(5): prof.step() with record_function("## forward ##"): pred = model(inputs)
with record_function("## backward ##"): loss_fn(pred, labels).backward()
with record_function("## optimizer ##"): optimizer.step() optimizer.zero_grad(set_to_none=True)
# Construct the memory timeline HTML plot. prof.export_memory_timeline(f"{file_prefix}.html", device="cuda:0")