Speeding Up Pytorch Models

Are you developing an AI App that should compute in Real-Time? Are throughput and latency your major concerns?

If you are developing AI software that should get results from an AI model in real-time or you are dealing with a huge stream of data and hence model’s speed is your major concern, and you are looking for a way to optimize the speed of your model, then Hey! You are reading the right blog right now. In this blog, we will use TensorRT to improve the latency of a Pytorch model with a significant margin. The actual code can be found here.

TensorRT

NVIDIA® TensorRT™ is an ecosystem of APIs for high-performance deep learning inference. TensorRT includes inference runtime and model optimizations that deliver low latency and high throughput for production applications. For more details check TensorRT

PyTorch

PyTorch is an open-source machine learning (ML) framework based on the Python programming language and the Torch library. Torch is an open-source ML library used for creating deep neural networks and is written in the Lua scripting language. It’s one of the preferred platforms for deep learning research. The framework is built to speed up the process between research prototyping and deployment. Check this for more on Pytorch.

We’ll be developing a Classification model on amnist dataset using PyTorch and Convolutional Neural Networks(CNNs).

We’ll start by importing the necessary modules, tensorRT, torch_tensorrt along with MNIST dataset.

Then we load the mnist dataset

This loads mnist dataset, which contains 60000 samples of handwritten numbers, we also convert these data samples to Tensors so they can be used efficiently on GPU with Pytorch.

We then split these data points into Train and Validation with 50000 going for training and 10000 for validation. Afterward we create a dataloader to effectively load these datapoints for training and validation of our torch model.

We will now define our Pytorch model for classification. This model includes three convolution layers, two dropouts for effective regularization, and two Linear layers.

Afterward we’ll define a classifier over this model, to effectively run training and validation of the model.

You can change hyperparameters defined in the classifier along with the input and output layers of Amnist model class. In fact I highly encourage you to play with multiple settings to see the effect of these hyperparameters on model accuracy and speed.

We then define train and validation functions for this class.

We now train this model on 10 epochs

So we see that after running 10 epochs, we get an accuracy of around 91% on the validation set.

Now we will define an inference function similar to train/val but we will note inference time to get fps or precisely datapoints per second of the model.

This function along with the accuracy prints average (datapoints per second) DPPS as FPS processed by our model.

You may wonder why I’m skipping the batch if its size is not 500. I’ll explain it when we compile our model with torch_tensorrt.

We now load the best checkpoint and run inference on it with val_loader.

We can see that the average throughput of the model is 3686 data points per second.

Let’s Make This Model FAST!!!!

First, we will convert this model to torch script module with torch.jit.script.

Scripting an nn.Module by default will compile the forward method and recursively compile any methods, submodules, and functions called by forward. If a nn.Module only uses features supported in TorchScript, no changes to the original module code should be necessary. script will construct [ScriptModule](https://pytorch.org/docs/stable/generated/torch.jit.ScriptModule.html#torch.jit.ScriptModule) that has copies of the attributes, parameters, and methods of the original module.

We now compile our model with torch_tensorrt. This is the magic step.

Here we are providing script module to compile the function of torch_tensorrt, along with an example input size. Input size is necessary with consideration that any future input size for the compiled model will be the same, otherwise, this model won’t work.

dtype specify the datatype of your input tensor, enabled_precisions allows us to quantize model to lower precisions like float16 or int8, and workspace size takes vram size of your GPU.

Let’s run inference on the compiled model now.

Wow!!! Our new model is 4.86 times faster than the actual torch model while accuracy is same as the original model. Lucky for this case, you can experience a drop of 2 or 3 percent sometimes depending on the complexity of the model. You can play with other parameters of torch tensorrt if you can tradeoff some accuracy with latency.

Conclusion

In today's blog, we wrote a simple code to optimize the latency of a torch model. We gained around 4.8x boost on our classification model without losing any accuracy. One thing to notice is that this approach works on Linux environments only since torch tensorrt is not yet supported for Windows. There is one complex way to optimize the model for Windows through C++, LibTorch, and TensorRT package for C++.