Best Performance Tuning Practices for Pytorch
By: Rishik C. Mourya
Tags: Efficiency, Benchmark, Pytorch, Performance, Optimization

A rticle Overview

Alright, so Pytorch has been one of my de facto choices for handling complex deep learning pipelines, like reading complex data, creating heavy models, and how can I forget the debugging. Now Pytorch already is pretty fast and memory-efficient, however, I will not hesitate to make my script even faster :) With that being said, let’s have a look at some of the very simple guidelines, that can easily be applied and make your model train both faster and efficiently.

Enabling Asynchronous Data Loading And Preprocessing

In case you don’t know, Pytorch does support asynchronous data loading, that is reading, preprocess, and applying the data augmentations asynchronously in the separate worker processes. However, by default, all that is done synchronously, so all the required computations are done in the main thread, and hence blocking everything else.

One of the most heaviest operation you can do to slow down your script is, transferring a lot of tensors from CPU to GPU and vice-versa. By a lot of tensors, I don’t mean the size of the tensors, but how frequently you make this transfer. However, by default, Pytorch does not use pinned memory, which means this CPU to GPU mem copies would be synchronous as well.

Let’s solve both of these issues.

Using more than zero workers

You can simply set the argument for number of workers greater than 0. This argument assigns how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process.

torch.utils.data.DataLoader(dataset, batch_size, shuffle, num_workers = 4)

Note, you cannot just set this argument anything. Getting the right value for num_workers depends on a lot of factors. Setting the value too high could cause quite a lot of issues such as:

Might increase the memory usage and that is the most serious overhead.

Might cause high IO usage which can ultimately become very ineffective.

If the value is higher then the number of cores you have available, then most probably you’ll block everything out, and its highly likely that your whole program will freeze. (Based on my observations)

And there is not any thumb rule that you can follow, some people suggests to set the value to number of cpu cores you have-1, subtracting one to make sure that you script never breaks. But this still is not an optimal solution. The optimal value depends on the number of cpu cores you have available, the batch size for your dataloader and amount of IO usage which is required to load a single batch and of course your model.

For a given batch size, the best practice is to increase the num_workers slowly and stop once you see no more improvement in your training speed.

If possible, you can also try experimenting different values for batch size and num_workers.

Experiment results for different sets of batch size and num_workers. Source

For instance, in the above experiment you can see that the most performance boost for all different batch sizes was in case of setting num_workers to 1, rest of the values didn’t show any performance improvements, in fact setting the num_workers to available cpu cores was causing the most performance regression.

Using pinned memories

If you read your data samples in the CPU, and during training would like to transfer the batch to the GPU, then you can speed up this device transfer by setting pin_memory to True.

This lets your torch.utils.data.DataLoader allocate the data samples in page-locked memory, and therefore speeding up the transfer. Host to GPU copies are much faster when they originate from pinned (page-locked) memory.

You can set pin memory to True by passing this as an argument in DataLoader:

torch.utils.data.DataLoader(dataset, batch_size, shuffle, pin_memory = True)

It is always okay to set pin_memory to True for the example I explained above. Though when your dataset is so small, that you can simply put it to the GPU prior to the training, then pin_memory wouldn’t work of course.

Enable cuDNN Autotuner

Nvidia’s cuDNN supports many algorithms to compute the output of a CNN layer for a fixed size. So if you’re working with CNNs then its highly recommended to enable the cuDNN’s autotuner.

When this is enabled, then for the very first forward pass to the model, the cuDNN would run a short benchmark and try to find the best algorithm to optimize the computation cost for your given hardware and fixed input size, and after that first pass, your rest of the predictions would be much more faster and efficient.

Once again, by default this is disabled in Pytorch and can be enabled as follows

torch.backends.cudnn.benchmark = True

Here, the important thing to note is when this is enabled, then the input size should be fixed for making predictions. Because cuDNN will benchmark every time a new size appears, and most likely leading to worse runtime performances.

Now it’s not like that in your dataset the number of samples is perfectly divisible by the specified batch size, by default the last batch given by DataLoader would have number of samples less than the specified batch size.

So to solve this issue, you can simply pass drop_last argument in the DataLoader as follows:

torch.utils.data.DataLoader(dataset, batch_size, shuffle, drop_last = True)

This will make the DataLoader drop (ignore) the last batch with size less than the specified batch size, hence making the cuDNN autotuner works as expected. And depending on your hardware and model, you could get performance improvement of the range 1.2 to 1.7 times.

Don’t Optimize Worthless Parameters

Almost every self taught ML developer has watched Andrew Ng’s courses, where he explained why using bias parameters in a layer that is followed by batch normalization is redundant. Though this still is one the most common mistake I’ve seen.

Let me summarize what he explained.

When you apply batch normalization, you first normalize the layer by subtracting the mean and dividing by the standard deviation (where both are computed over the dimension of the batch).

Now since, the bias parameters are used in a simple addition operation, and the value of bias parameters remains same across the dimension of the batch, thus while subtracting the batch mean, you’d basically be subtracting the bias parameters you’ve just added (since the mean of let’s say k numbers, where each number is b, is indeed b), hence making the bias addition operation obsolete.

So next time, don’t use bias parameters in a layer which is followed by batch normalization. In Pytorch you can disable bias parameters by simply passing bias = False in any layer that has bias as one of its parameters.

Note, the above explanation is only true when you’re applying batch normalization before the activation layer. If an activation function is applied to the affine transformation before the batch normalization, then bias parameters are not wasted.

So should we apply batch normalization before or after the activation layer? That’s the topic of another article I guess.

Using Large Batch Size

Now this is pretty straight forward. It is recommended to increase the batch size as much as possible, this helps to minimize the cost of launching kernels in GPU, and also since many operations would be done parallel, so yeah… it would be pretty damn fast. In fact, doubling batch size would train your model faster than any other techniques I’m explaining in this article when applied alone.

But, the increasing batch size would lead to two major problems:

Not significant loss convergence.

And running out of GPU memory.

The first issue is very classic, and you can’t do much about it, unless you change your model accordingly. But my suggestion is to train your model with batch size such that during training you’re almost completely filling the GPU memory, then after few epochs when the loss is not further decreasing, then decrease the batch size, simple as that.

Also, try using learning rate warm-ups and weight decaying, these techniques also help the loss convergence for the most part.

And this way you’d be able to get the best from both worlds :)

Now, regarding the out of memory issue. I’ve written an awesome article, that solves this almost all the times, so definitely check it out.

Basically in Pytorch, you can use AMP (automatic mixed precision) that makes both forward and backward pass way faster and efficient, which allows to train the model much faster with high efficiency, thus less memory consumption.

Zeroing The Gradients Efficiently

This particular technique was contributed to Pytorch by Nvidia engineers. Have a look.

It’s highly common practice to zero out gradients before computing them for the next batch as follows:

model.zero_grad()

# ======== OR ========

optimizer.zero_grad()

Nvidia mentioned that this is not an optimal way to zero the gradients out. Surprised? Hell yeah me too, cuz in all the official tutorials and documentation of Pytorch, this is exactly how its done, and so does everyone.

They mentioned that the problems with this approach are:

These zero_grad methods will execute the memset to zero in a loop over all parameters. And this loop will call separate cuda kernels for every parameter.

Under the hood for backward pass, the Pytorch will use “plus equals” operator (“+=”) update the gradients ( because Pytorch by default adds the current gradients to the previous ones, this is the reason why we need to zero out the gradients for each batch in the first place ), and these operators first performs the read operation to obtain the original values and then the addition values.

), and these operators first performs the read operation to obtain the original values and then the addition values. After all above operations, the parameters are once again stored back in the memory.

What they instead suggested is to set the grad attribute of each parameter to none as follows:

for param in model.parameters():

param.grad = None

This approach solves the above issues:

Does not execute memset for each of the parameters.

The memory is instead zero out in more efficient manner within the Pytorch backend.

And the update of the gradients during backward pass will be done directly by simply using the equals operator (“=”), so no time is wasted in any unnecessary reading.

Conclusion

Data preprocessing and augmentation can be done asynchronously by using more than zero number of workers.

Faster device transfer can be achieved by pinned memory.

Models using CNN layers can be further tuned and optimized using cuDNN’s autotuner.

Setting a large batch size can obviously provide a whole lot of performance boost.

AMP can be used for more efficient and faster training.

And finally, zeroing out the gradients by setting the grad attribute to None solves a lot of performance overhead issues.

Alrighty then, I hope these techniques would be helpful for ya.

Thanks for reading.