@ariG23498: There are two official torch tutorials that link to the "Profiling in PyTorch" sereis! https://docs.pytorch.org/tutoria…
Summary
Two official PyTorch tutorials linking to the 'Profiling in PyTorch' series are highlighted, providing guidance on using the PyTorch profiler API for performance debugging.
View Cached Full Text
Cached at: 07/20/26, 05:29 PM
There are two official torch tutorials that link to the “Profiling in PyTorch” sereis! https://docs.pytorch.org/tutorials/beginner/profiler.html… https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html…
Profiling your PyTorch Module
Source: https://docs.pytorch.org/tutorials/beginner/profiler.html Note
Go to the endto download the full example code.
Created On: Dec 30, 2020 | Last Updated: Jul 19, 2026 | Last Verified: Nov 05, 2024
Author:Suraj Subramanian
PyTorch includes a profiler API that is useful to identify the time and memory costs of various PyTorch operations in your code. Profiler can be easily integrated in your code, and the results can be printed as a table or returned in a JSON trace file.
Note
Profiler supports multithreaded models. Profiler runs in the same thread as the operation but it will also profile child operators that might run in another thread. Concurrently-running profilers will be scoped to their own thread to prevent mixing of results.
Note
PyTorch 1.8 introduces the new API that will replace the older profiler API in the future releases. Check the new API atthis page.
Head on over tothis recipefor a quicker walkthrough of Profiler API usage.
import torch
import numpy as np
from torch import nn
import torch.autograd.profiler as profiler
Performance debugging using Profiler#
Profiler can be useful to identify performance bottlenecks in your models. In this example, we build a custom module that performs two sub-tasks:
- a linear transformation on the input, and
- use the transformation result to get indices on a mask tensor.
We wrap the code for each sub-task in separate labelled context managers usingprofiler\.record\_function\("label"\). In the profiler output, the aggregate performance metrics of all operations in the sub-task will show up under its corresponding label.
Note that using Profiler incurs some overhead, and is best used only for investigating code. Remember to remove it if you are benchmarking runtimes.
class MyModule(nn.Module):
def __init__(self, in_features: int, out_features: int, bias: bool = True):
super(MyModule, self).__init__()
self.linear = nn.Linear(in_features, out_features, bias)
def forward(self, input, mask):
with profiler.record_function("LINEAR PASS"):
out = self.linear(input)
with profiler.record_function("MASK INDICES"):
threshold = out.sum(axis=1).mean().item()
hi_idx = np.argwhere(mask.cpu().numpy() > threshold)
hi_idx = torch.from_numpy(hi_idx).cuda()
return out, hi_idx
Profile the forward pass#
We initialize random input and mask tensors, and the model.
Before we run the profiler, we warm-up CUDA to ensure accurate performance benchmarking. We wrap the forward pass of our module in theprofiler\.profilecontext manager. Thewith\_stack=Trueparameter appends the file and line number of the operation in the trace.
Warning
with\_stack=Trueincurs an additional overhead, and is better suited for investigating code. Remember to remove it if you are benchmarking performance.
model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.double).cuda()
# warm-up
model(input, mask)
with profiler.profile(with_stack=True, profile_memory=True) as prof:
out, idx = model(input, mask)
Print profiler results#
Finally, we print the profiler results.profiler\.key\_averagesaggregates the results by operator name, and optionally by input shapes and/or stack trace events. Grouping by input shapes is useful to identify which tensor shapes are utilized by the model.
Here, we usegroup\_by\_stack\_n=5which aggregates runtimes by the operation and its traceback (truncated to the most recent 5 events), and display the events in the order they are registered. The table can also be sorted by passing asort\_byargument (refer to thedocsfor valid sorting keys).
Note
When running profiler in a notebook, you might see entries like<ipython\-input\-18\-193a910735e8\>\(13\): forwardinstead of filenames in the stacktrace. These correspond to<notebook\-cell\>\(line number\): calling\-function.
print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))
"""
(Some columns are omitted)
------------- ------------ ------------ ------------ ---------------------------------
Name Self CPU % Self CPU Self CPU Mem Source Location
------------- ------------ ------------ ------------ ---------------------------------
MASK INDICES 87.88% 5.212s -953.67 Mb /mnt/xarfuse/.../torch/au
<ipython-input-...>(10): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
aten::copy_ 12.07% 715.848ms 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
LINEAR PASS 0.01% 350.151us -20 b /mnt/xarfuse/.../torch/au
<ipython-input-...>(7): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
aten::addmm 0.00% 293.342us 0 b /mnt/xarfuse/.../torch/nn
/mnt/xarfuse/.../torch/nn
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(8): forward
/mnt/xarfuse/.../torch/nn
aten::mean 0.00% 235.095us 0 b <ipython-input-...>(11): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
----------------------------- ------------ ---------- ----------------------------------
Self CPU time total: 5.931s
"""
Improve memory performance#
Note that the most expensive operations - in terms of memory and time - are atforward \(10\)representing the operations within MASK INDICES. Let’s try to tackle the memory consumption first. We can see that the\.to\(\)operation at line 12 consumes 953.67 Mb. This operation copiesmaskto the CPU.maskis initialized with atorch\.doubledatatype. Can we reduce the memory footprint by casting it totorch\.floatinstead?
model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.float).cuda()
# warm-up
model(input, mask)
with profiler.profile(with_stack=True, profile_memory=True) as prof:
out, idx = model(input, mask)
print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))
"""
(Some columns are omitted)
----------------- ------------ ------------ ------------ --------------------------------
Name Self CPU % Self CPU Self CPU Mem Source Location
----------------- ------------ ------------ ------------ --------------------------------
MASK INDICES 93.61% 5.006s -476.84 Mb /mnt/xarfuse/.../torch/au
<ipython-input-...>(10): forward
/mnt/xarfuse/ /torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
aten::copy_ 6.34% 338.759ms 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
aten::as_strided 0.01% 281.808us 0 b <ipython-input-...>(11): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
aten::addmm 0.01% 275.721us 0 b /mnt/xarfuse/.../torch/nn
/mnt/xarfuse/.../torch/nn
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(8): forward
/mnt/xarfuse/.../torch/nn
aten::_local 0.01% 268.650us 0 b <ipython-input-...>(11): forward
_scalar_dense /mnt/xarfuse/.../torch/nn
<ipython-input-...>(9): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
----------------- ------------ ------------ ------------ --------------------------------
Self CPU time total: 5.347s
"""
The CPU memory footprint for this operation has halved.
Improve time performance#
While the time consumed has also reduced a bit, it’s still too high. Turns out copying a matrix from CUDA to CPU is pretty expensive! Theaten::copy\_operator inforward \(12\)copiesmaskto CPU so that it can use the NumPyargwherefunction.aten::copy\_atforward\(13\)copies the array back to CUDA as a tensor. We could eliminate both of these if we use atorchfunctionnonzero\(\)here instead.
class MyModule(nn.Module):
def __init__(self, in_features: int, out_features: int, bias: bool = True):
super(MyModule, self).__init__()
self.linear = nn.Linear(in_features, out_features, bias)
def forward(self, input, mask):
with profiler.record_function("LINEAR PASS"):
out = self.linear(input)
with profiler.record_function("MASK INDICES"):
threshold = out.sum(axis=1).mean()
hi_idx = (mask > threshold).nonzero(as_tuple=True)
return out, hi_idx
model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.float).cuda()
# warm-up
model(input, mask)
with profiler.profile(with_stack=True, profile_memory=True) as prof:
out, idx = model(input, mask)
print(prof.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=5))
"""
(Some columns are omitted)
-------------- ------------ ------------ ------------ ---------------------------------
Name Self CPU % Self CPU Self CPU Mem Source Location
-------------- ------------ ------------ ------------ ---------------------------------
aten::gt 57.17% 129.089ms 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(25): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
aten::nonzero 37.38% 84.402ms 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(25): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
INDEX SCORE 3.32% 7.491ms -119.21 Mb /mnt/xarfuse/.../torch/au
<ipython-input-...>(10): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(25): <module>
/mnt/xarfuse/.../IPython/
aten::as_strided 0.20% 441.587us 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(25): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
aten::nonzero
_numpy 0.18% 395.602us 0 b <ipython-input-...>(12): forward
/mnt/xarfuse/.../torch/nn
<ipython-input-...>(25): <module>
/mnt/xarfuse/.../IPython/
/mnt/xarfuse/.../IPython/
-------------- ------------ ------------ ------------ ---------------------------------
Self CPU time total: 225.801ms
"""
Further Reading#
We have seen how Profiler can be used to investigate time and memory bottlenecks in PyTorch models. Read more about Profiler here:
- Profiler Usage Recipe
- Profiling RPC-Based Workloads
- Profiler API Docs
- Hugging Face PyTorch Profiler Blog Post
Aritra 🤗 (@ariG23498): The Profiling Blog Post is now listed in the official torch profiler documentation! 🤗
Similar Articles
@ariG23498: This man can write! Pair this understanding with profiling (shameless plug to my series) and you are golden! http://hf.…
A beginner-friendly guide to using torch.profiler in PyTorch for profiling and optimizing deep learning workloads, covering trace reading, CUDA analysis, and torch.compile integration.
Profiling in PyTorch (Part 1): A Beginner's Guide to torch.profiler
A beginner-friendly guide to using PyTorch's torch.profiler for profiling and optimizing neural network operations, starting with matrix multiplication and bias addition. It explains how to read profiler traces and understand CPU/GPU interactions.
@ariG23498: It is profiling time! In Part 2 we cover: > trace a Linear layer > talk about mul + add vs linear > gemm epilogues (my …
Announces Part 2 of a profiling tutorial covering linear layer tracing, gemm epilogues, MLP tracing, and comparisons of torch compile vs Liger kernels, with a link to the full content.
@RisingSayak: I realized that what I cannot profile, I cannot optimize. This is why I embarked on a little project in Diffusers, to t…
Sayak Paul describes a project to profile and optimize Diffusers pipelines using torch.compile, and announces a tutorial series by Ari G. on the topic.
Profiling in PyTorch (Part 3): Attention is all you profile
This tutorial demonstrates how to profile various attention implementations in PyTorch, from naive attention to scaled dot-product attention with different backends, using the PyTorch profiler on an NVIDIA A100 GPU.