Train Neural Networks without Draining your Pocket: Understand PyTorch Lightning’s Profiler Table

Use Lightning’s profiler to interpret if your model is too complex for the data and hardware. Learn to read the logs and identify the bottlenecks.

I’m guessing you are here because,

• You’ve gone through Profile your PyTorch Models and you’re here to figure out what to even do with that complicated profiler output
• Or you’ve tried out Lightning’s profiler by yourself and you’re unable to understand the output

Fret not, by the end of this blog, at least some of your concerns should be resolved. Yes, I did say some. Profiling is a massive topic by itself and given the pace at which models are getting heavier, this topic is just going get more important by the day!

Let’s use the profile from the previous part as reference.

UNDERSTANDING PROFILER LOGS

This was the output of the profiler code (Kaggle notebook, GitHub). Based on the platform you are using you may have to zoom on the image to read it! We’ll call this output the Profiler Table and use it to understand where the hardware (both CPU and GPU) spends most of their time during the training and validation process.

Since we had set with_modules=True in the code, notice how specific details like resnet.BasicBlock are logged along with the raw math operations. I know, its still a little difficult to read but it could be way worse!

EXPLORING THE PROFILE LOGS: COLUMNS

The columns as you may have noticed already is all about the hardware. What you’re looking at his how much time the hardware spends on each stage in the training in terms of percentage and milliseconds (occasionally microseconds).

Self CPU(% and time): Displays the time spent only on the specific function (in the row), excluding calls to other sub-functions.

CPU total(% and time): Shows the total time spent on that function (in the row) including all the sub-functions it called.

Note: The self and total situation may sound confusing but we will discuss once we discuss all the rows and columns!

Self CUDA(% and time): Displays the time your GPU spent executing the specific function (in the row). This is a column where your bottlenecks will often appear!

CUDA total: This displays the total GPU time for this function and everything inside it.

# of Calls: How many times this particular operation was executed during the profiling window.

CPU time average: This is the mean latency of the CPU-side instruction, which shows the mean of the CPU total for a single execution of the function (i.e., CPU total/# of calls)

CUDA time average: This is the mean latency of the GPU kernel execution which indicates the mean latency of the GPU kernel execution (i.e., CUDA total/# of calls)

Understanding Self vs Total (CPU or CUDA)

The difference between Self (CPU/CUDA) and Total (CPU/CUDA) is the difference between local execution and call-stack accumulation. Look a basic block from the ResNet18 architecture — two Convolution layers, two Batch Normalization layers, a ReLU activation and the skip connection.

Self CPU for the basic ResNet block represents the time the CPU spent executing the internal logic of the BasicBlock class, excluding any time spent in the layers it contains. This means it counts the overhead of the Pythin forward() method, the addition operation for the shortcut (i.e., x+F(x) as shown in the figure above) and any internal variable assignments. It ignores the parts where the code enters a nn.Conv2d or nn.Batch2D layer.

A high value for Self CPU for the basic block would usually indicate that the skip connection is adding an overhead.

Total CPU indicates the entire duration spent from the basic block’s forward() is called till the output of the block is computed. That means, this includes the skip connection along with the time spent in the two Convolution, Batch normalization layers, and the ReLU layer.

A low Self CPU time but high Total CPU time would indicate that the bottleneck is possibly the convolution operations in the block.

The same explanations can be extrapolated to Self CUDA and Total CUDA!

EXPLORING THE PROFILE LOGS: ROWS

The rows denote the names of the functions/operations that are running (and being profiled). Depending on which profiler you use, the rows can be slightly tricky to understand. We will use our current profiler for reference but keep in mind the rows will look very different based on the selected profiler.

Notice how the profiler rows have these details inside []. Check the image below where I’ve highlighted a few of these,

These are known as Context tags or Metadata labels. These are hierarchical identifiers that map low-level hardware execution to high-level software abstraction. Think of it as a way to connect the model architecture to the hardware which is an important step for identifying model bottlenecks. Lets look at each of these in detail,

• High-Level Framework Names ([pl])

These rows represent the Lightning Execution Engine that manages the high-level training loop logic. Some the terms you will notice in the profile are,

*ProfilerStep: This is the top-level container for one iteration of our profiling schedule

SingleDeviceStrategy.validation_step: This confirms that we are looking at the validation phase and shows the total time to process a batch and compute the metrics.

torchmetrics.classification.accuracy: This denotes the time consumed in computing the accuracy metric.

You might notice another tag as ([strategy]). This can be considered as a sub-component under this heading and is used to denote the training strategy (ex SingleDeviceStrategy, DDP etc.). The strategy is an indicator of how the data flows between the CPU and GPU and can help you figure out if the bottleneck is in the data flow rather than the specifics of the model.

• Neural Network Modules ([module])

Rows with module define the specific PyTorch nn.Module (ex. ResNet or Sequential) defined in our code. Some of the terms you’ll notice in the profile are,

torchvision.models.resnet.BasicBlock: In our case, these denote the specific layers of the ResNet18 backbone. If they appear in the profile, it means they are currently active (not frozen) and being used in the forward pass. The total time for this will show the cost of processing one full residual block.

torch.nn.modules.container.Sequential: This represents the self.feature_extractor and self.classifier sections of the model.

• Native Operator Rows (aten::)

These are the ATen (A Tensor Library) functions that actually perform the linear algebra. They have no sub functions. In our profile, you will see these,

aten::conv2d / aten::convolution: This is the CPU-side call that prepares the instruction for the GPU.

aten::cudnn_convolution: This is the NVIDIA-Optimized Kernel, the highest-performance version of a convolution. The Self CUDA time for this primarily indicates how much heavy-lifting the GPU is doing for our backbone.

aten::batch_norm: This is the feature normalization step and it calculates the mean and variance of the feature maps to stabilize training.

aten::_batch_norm_impl_index: This is the low-level Implementation of batch normalization in C++ backend.

maxwell_ sgemm_32x128_n: This is a specific NVIDIA GPU kernel name that stands for Single-precision General Matrix Multiply. This is the math that powers the linear layers. 32x128 denotes the shapes here.

5x_cudnn_maxwell_scudnn_winograd: This denotes a Winograd Convolution kernel, which is a mathematical optimization that reduces the number of multiplications needed for a convolution, making it faster.

Now that we have all this information, let’s explore how to put all the information together and optimize our models in the next part!

REFERENCES

1. Lightning Documentation: https://lightning.ai/docs/pytorch/stable/tuning/profiler.html
2. PyTorch profiler: https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html