How to run the PyTorch profiler

PyTorch profiler captures the operators, Python ranges, and device activity that make up a training or inference step. Running it around a short representative loop shows whether time is going into matrix multiplies, loss calculation, dataloader work, or custom code labels before heavier optimization work begins.

The torch.profiler.profile context manager records selected activity groups. A schedule keeps warm-up overhead out of the saved result, and profiler.step() marks iteration boundaries so the active window lines up with model steps. The portable smoke path records CPU activity and writes a TensorBoard trace directory with tensorboard_trace_handler().

Profiler runs add overhead, so use a short workload that keeps the same model shapes and batch path as the bottleneck being investigated. The text table from key_averages() is enough for a first operator ranking, while the trace file can be opened later in TensorBoard when timeline detail is needed.

1. Create a profiler smoke script around a small training step.

   profile_train_step.py

   from pathlib import Path
   import shutil
    
   import torch
   from torch import nn
   from torch.profiler import ProfilerActivity, profile, record_function, schedule, tensorboard_trace_handler
    
    
   torch.manual_seed(23)
    
   trace_dir = Path("profiler-traces")
   if trace_dir.exists():
       shutil.rmtree(trace_dir)
    
    
   class TinyClassifier(nn.Module):
       def __init__(self):
           super().__init__()
           self.layers = nn.Sequential(
               nn.Linear(8, 16),
               nn.ReLU(),
               nn.Linear(16, 3),
           )
    
       def forward(self, inputs):
           return self.layers(inputs)
    
    
   model = TinyClassifier()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
   loss_fn = nn.CrossEntropyLoss()
    
    
   def train_step(step):
       inputs = torch.randn(12, 8)
       targets = torch.tensor([0, 1, 2, 1, 0, 2, 2, 1, 0, 1, 2, 0])
    
       optimizer.zero_grad(set_to_none=True)
    
       with record_function("forward_and_loss"):
           outputs = model(inputs)
           loss = loss_fn(outputs, targets)
    
       with record_function("backward_and_step"):
           loss.backward()
           optimizer.step()
    
       print(f"step={step} loss={loss.item():.4f}")
    
    
   with profile(
       activities=[ProfilerActivity.CPU],
       schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
       on_trace_ready=tensorboard_trace_handler(str(trace_dir), worker_name="profiler-demo"),
       record_shapes=True,
       acc_events=True,
   ) as profiler:
       for step in range(5):
           train_step(step)
           profiler.step()
    
   trace_files = sorted(trace_dir.glob("*.pt.trace.json"))
    
   print(f"torch_version={torch.__version__}")
   print(f"trace_dir={trace_dir}")
   print(f"trace_file_count={len(trace_files)}")
   for path in trace_files:
       print(f"trace_file={path.name}")
    
   print(
       profiler.key_averages().table(
           sort_by="self_cpu_time_total",
           row_limit=6,
       )
   )

   ProfilerActivity.CPU keeps the smoke run portable. Add ProfilerActivity.CUDA only after the model and input tensors run on CUDA.
   Related: How to enable CUDA in PyTorch

2. Run the profiler smoke script.

   $ python profile_train_step.py
   step=0 loss=1.0979
   step=1 loss=1.1667
   step=2 loss=1.1166
   step=3 loss=1.1557
   step=4 loss=1.1061
   torch_version=2.12.1+cpu
   trace_dir=profiler-traces
   trace_file_count=1
   trace_file=profiler-demo.1783379189188587923.pt.trace.json
   -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                      Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
   -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                               aten::addmm        31.32%       2.663ms        31.72%       2.697ms     449.541us             6  
                                   aten::nll_loss_backward        14.36%       1.221ms        14.38%       1.223ms     407.667us             3  
                          aten::_log_softmax_backward_data        10.24%     870.916us        10.24%     870.916us     290.305us             3  
                                        aten::_log_softmax        10.10%     859.001us        10.10%     859.001us     286.334us             3  
                                          forward_and_loss         8.12%     690.377us        52.86%       4.495ms       1.498ms             3  
                                         backward_and_step         6.08%     516.672us        42.21%       3.589ms       1.196ms             3  
   -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
   Self CPU time total: 8.503ms

   Some PyTorch builds print profiler_start and profiler_stop diagnostics to stderr. The trace count and operator table are the success signals.

3. Check the top self-CPU rows before changing model code.

   In the smoke run, aten::addmm comes from the two linear layers. The forward_and_loss and backward_and_step rows come from the custom record_function() labels.

4. Move the profiler context into the real training loop after the smoke run works.

   with profile(
       activities=[ProfilerActivity.CPU],
       schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
       on_trace_ready=tensorboard_trace_handler("profiler-traces", worker_name="train-worker-0"),
       record_shapes=True,
       acc_events=True,
   ) as profiler:
       for step, (inputs, targets) in enumerate(train_loader):
           loss = train_one_batch(model, inputs, targets)
           profiler.step()
           if step >= 4:
               break

   Keep the active window short enough to capture representative batches without profiling an entire long training run. Use the trace directory as the TensorBoard logdir when timeline view is needed.
   Related: How to log PyTorch training metrics to TensorBoard

5. Remove the smoke script and generated trace directory after moving the pattern into the project.

   $ rm -r profile_train_step.py profiler-traces