PyTorch JIT and TorchScript. A path to production for PyTorch models | by Abhishek Sharma | Towards Data Science

We all love the dynamicity and ease of use of PyTorch. But when it comes to deployment these qualities are not as desirable as performance and portability.
In the past deployment required moving to a different framework. The move from research to production required 2 separate teams with different skill sets to make this possible. This process changed with the introduction of TorchScript.
With TorchScript, PyTorch aims to create a unified framework from research to production. TorchScript will take your PyTorch modules as input and convert them into a production-friendly format.
It will run your models faster and independent of the Python runtime.
PyTorch supports 2 separate modes to handle research and production environment.
First is the Eager mode. It is built for faster prototyping, training, and experimentation.
Second is the Script mode. It is focused on the production use case. It has 2 components PyTorch JIT and TorchScript.
In one line, it gets rids of Python’s GIL and dependence on Python runtime. A nuanced explanation is as follows
PortabilityPortability allows models to be deployed in multithreaded inference servers, mobiles, and cars which is difficult with Python. In order to achieve this PyTorch models needs to be decoupled from any specific runtime.PerformancePyTorch JIT is an optimizing JIT compiler for PyTorch. It uses runtime information to optimize TorchScript modules. It can automate optimizations like layer fusion, quantization, sparsification.
PortabilityPortability allows models to be deployed in multithreaded inference servers, mobiles, and cars which is difficult with Python. In order to achieve this PyTorch models needs to be decoupled from any specific runtime.
PerformancePyTorch JIT is an optimizing JIT compiler for PyTorch. It uses runtime information to optimize TorchScript modules. It can automate optimizations like layer fusion, quantization, sparsification.
Script mode creates an intermediate representation (IR) of your PyTorch Eager module(throughtorch.jit.trace/torch.jit.script). The IR is internally optimized and utilizes PyTorch JIT compilation at runtime. PyTorch JIT compiler uses runtime information to optimize the IR. This IR is decoupled from the Python runtime.
The script mode works by utilizing PyTorch JIT and TorchScript.
PyTorch JIT is an optimized compiler for PyTorch programs.
It is a lightweight threadsafe interpreterSupports easy to write custom transformationsIt’s not just for inference as it has auto diff support
It is a lightweight threadsafe interpreter
Supports easy to write custom transformations
It’s not just for inference as it has auto diff support
TorchScript is a static high-performance subset of Python language, specialized for ML applications. It supports
Complex control flowsCommon data structuresUser-defined classes
Complex control flows
Common data structures
User-defined classes
We have introduced PyTorch JIT, TorchScript, and Script mode. We have discussed, What are they, and why we need them? In the following segments, we will understand how to use them?
Script mode is invoked by either torch.jit.trace or torch.jit.script.
torch.jit.trace take a data instance and your trained eager module as input. The tracer runs the supplied module and records the tensor operations performed. This recording is turned into a TorchScript module.
Skip to Example 1 to see it in action
It can reuse existing eager model code and can handle almost any program with exclusive torch tensors/operations.
Its main drawback is that it omits all control flow, data structures, and python constructs. It can also create unfaithful representations without any warnings. Always checks its IR to see if it is correctly parsing the structure of your PyTorch model.
torch.jit.script allows you to write your code directly into TorchScript. It's more verbose but it more versatile and with a little tweaking can support the majority of the PyTorch models.
In contrast to trace mode you only need to pass an instance of your model/module to torch.jit.script. A data sample is not necessary.
Skip to Example 2 on its practical utility
It preserves control flows and other python constructs and looks more like Python. It has first-class support for lists/dictionaries.
It omits constant values and requires typecasting. If no type is provided it defaults to tensor
That’s too much theory for one sitting, Let’s move over to the code and see the trace/script in action.
BERT (Bidirectional Encoder Representations from Transformers) was developed by researchers at Google AI. In the following example, we utilize BERT from the transformer’s library provided by HuggingFace.
This example is divided into 3 parts
Part 1Initializes BERT model/tokenizers and creates a sample data for inference
Part 2:Prepares PyTorch models for inference on CPU/GPU.
Model/Data should be on the same device for training/inference to happen. cuda() transfers the model/data from CPU to GPU.
Step 3:Prepares TorchScript modules for inference on CPU/GPU
On CPU the runtimes are similar but on GPU TorchScript clearly outperforms PyTorch.
TorchScript creates an IR of the PyTorch models which can be compiled optimally at runtime by PyTorch JIT. This IR can be viewed using traced_model.code.
ResNet, short for Residual Networks is a classic neural network used as a backbone for many computer vision tasks. This model was the winner of the ImageNet challenge in 2015.
The following example has 2 segments
Example 1.1 / 1.2:PyTorch ResNet initialization and inference on CPU/GPU
Example 2.1/2.2:torch.jit.script initialization and inference on CPU/GPU
TorchScript significantly outperforms the PyTorch implementation on GPU.
TorchScript saves/loads modules into an archive format. This archive is a standalone representation of the model and can be loaded into an entirely separate process.
Saving a module torch.jit.save(traced_model,’traced_bert.pt’)
Loading a module loaded = torch.jit.load('traced_bert.pt')
Example 1: Use traced_model.code to view PyTorch IRSkipping this as it’s very verbose
Example 2: Use script_cell_gpu.code to view PyTorch IR
def forward(self,    input: Tensor) -> Tensor:  _0 = self.fc  _1 = self.avgpool  _2 = self.layer4  _3 = self.layer3  _4 = self.layer2  _5 = self.layer1  _6 = self.maxpool  _7 = self.relu  _8 = (self.bn1).forward((self.conv1).forward(input, ), )  _9 = (_5).forward((_6).forward((_7).forward(_8, ), ), )  _10 = (_2).forward((_3).forward((_4).forward(_9, ), ), )  input0 = torch.flatten((_1).forward(_10, ), 1, -1)  return (_0).forward(input0, )
torch.jit.trace and torch.jit.script can be combined to cover for their deficiencies. Look at this official example. They do this by inlining the code.
Apart from improved inference time, there are other benefits of using TorchScript
TorchScript decouples your model from any runtime environment. It gets rid of Python’s GIL which is a major bottleneck for multithreaded inference.
TorchScript focuses on whole program optimizations.
TorchScript automatically optimizes common patterns in neural networks to improve latency and throughput.
TorchScript modules are exportable to a wide variety of environments, from C++ servers to mobile.
This blog has introduced script mode. Functionally the script mode provides tools to take your research models to production while still being in the PyTorch ecosystem.
To reiterate this is how a typical path for model development and deployment would look like.
Step 0: Develop your models, preferably in the PyTorch Eager mode
Step 1: Create TorchScript module by using eithertorch.jit.trace or/andtorch.jit.script on your PyTorch model
Step 2: Transfer these modules to the production environment using torch.jit.save/torch.jit.load. In this format, they can be run anywhere from servers to edge devices
Step 3: Use the PyTorch JIT compiler to optimize these programs at inference time and enjoy faster inference with minimal effort.
Torchscript + PyTorch JITResearch to ProductionPyTorch DocumentationUsing TorchScript for Transformer models
Torchscript + PyTorch JIT
Research to Production
PyTorch Documentation
Using TorchScript for Transformer models