Profiling PyTorch Code

This notebook demonstrates how to incorporate PyTorch Kineto's Tensorboard plugin for profiling PyTorch code with PyTorch Lightning as the high-level training API and Weights & Biases as the logging solution.

The final result looks something like what you see below:

The work done by processes, threads, and streams on the CPU and GPU is displayed along with precise timing information in an interactive viewer that can be incorporated into Weights & Biases workspaces and Reports or exported to external viewers.

That means you can incorporate tracing and profiling into your model training and evaluation pipeline -- storing, visualizing, and communicating performance results alongside other key metrics and metadata, like loss curves, hard examples from datasets, and hyperparameter optimization results.

NB: This tool is based on the Chrome Trace Viewer, which works best with that browser.

Set Up Profiled Training

Network Module

To profile neural network code, we first need to write it.

For this demo, we'll stick with a simple LeNet-style DNN, based on the PyTorch introductory tutorial.

To get this module to work with PyTorch Lightning, we need to define two more methods, which hook into the training loop.

Check out this tutorial video and notebook for more on using PyTorch Lightning and W&B.

Profiler Callback

The profiler operates a bit like a PyTorch optimizer: it has a .step method that we need to call to demarcate the code we're interested in profiling.

A single training step (forward and backward prop) is both the typical target of performance optimizations and already rich enough to more than fill out a profiling trace, so we want to call .step on each step.

The cell below defines a quick-and-dirty method for doing so in PyTorch Lightning using the Callback system.

Run Profiled Training

We're now ready to go!

The cell below creates a DataLoader based on the information in the configuration dictionary. Choices made here have a substantial impact on performance and show up very markedly in the trace.

After you've run with the default values, check out the creation of the trainloader and the trainer for comments on what these arguments do and then try a few different choices out, as suggested below.

Reading Profiling Results

Head to the Artifacts tab (identified by the "stacked pucks" database icon) for your W&B run page, at the URL that appears in the output of the cell above, then select the artifact named trace-. In the Files tab, select trace.pt.trace.json to pull up the Trace Viewer.

You can also check out an example from an earlier run here.

The trace shows which operations were running and when in each process/thread/stream on the CPU and on the GPU.

In the main thread (the one in which the Profiler Steps appear), locate the following steps:

1. the loading of data (hint: look for enumerate on the CPU, nothing on the GPU)
2. the forward pass to calculate the loss (hint: look for simultaneous activity on CPU+GPU, with aten in the operation names)
3. the backward pass to calculate the gradient of the loss (hint: look for simultaneous activity on CPU+GPU, with backward in the operation names).

If you ran with the default settings (in particular, num_workers=0), you'll notice that these steps are all run sequentially, meaning that between loading one batch and loading the next, the DataLoader is effectively idling, and during the loading of a batch, the GPU is idling.

Change num_workers in the config to 1 or 2 and then re-execute the cell above. You should notice a difference, in particular in the fraction of time the GPU is active. (Note: the DataLoader may even be hard to find in this case!)

For more on how to read these results, check out this W&B Report.