Skip to content

Commit

Permalink
add comment
Browse files Browse the repository at this point in the history
  • Loading branch information
rayg1234 committed Jul 10, 2024
1 parent 8e0e612 commit 3afd454
Showing 1 changed file with 20 additions and 1 deletion.
21 changes: 20 additions & 1 deletion src/fairchem/core/common/profiler_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,26 @@ def trace_handler(p):
return trace_handler

def get_profile_schedule(wait: int = 5, warmup: int = 5, active: int = 2):
"""Get a profile schedule"""
"""Get a profile schedule and total number of steps to run
check pytorch docs on the meaning of these paramters:
https://pytorch.org/docs/stable/profiler.html#torch.profiler.schedule
Example usage:
```
trace_handler = get_default_profiler_handler(run_id = self.config["cmd"]["timestamp_id"],
output_dir = self.config["cmd"]["results_dir"],
logger = self.logger)
profile_schedule, total_profile_steps = get_profile_schedule()
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=profile_schedule,
on_trace_ready=trace_handler
) as p:
for i in steps:
<code block to profile>
if i < total_profile_steps:
p.step()
"""
total_profile_steps = wait + warmup + active
profile_schedule = torch.profiler.schedule(wait=wait, warmup=warmup, active=active)

Expand Down

0 comments on commit 3afd454

Please sign in to comment.