Optimizing PyTorch Models: How TorchScript Works and How to Use It
Hello!
Today's topic is optimization, an essential step when taking an AI application developed with PyTorch to production. Specifically, we will look at optimizing various trained models using TorchScript.
TorchScript Basics
1 What Is TorchScript?
TorchScript is a technology that converts PyTorch models into an optimized intermediate representation (IR).
That may sound a little abstract, though.
To put it in plainer terms:
In short, it is a technology for running machine learning models built with PyTorch fast, and in a wide variety of environments.
For example, when you want to...
- Run your model in environments where Python is not installed
- Use it on smartphones and other embedded devices
- Dramatically speed up execution
- Run multiple operations concurrently and efficiently
...TorchScript is a great choice.
In other words, TorchScript is extremely useful when deploying models in production services.
2 How TorchScript Works
Let's walk through the basic mechanism of converting a PyTorch model to TorchScript.
2.1 Defining the Model
First, let's build a simple CNN model for image recognition
import torch
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(3, 64, 3)
self.relu = torch.nn.ReLU()
self.pool = torch.nn.MaxPool2d(2)
# Compute the input dimension of the fully connected layer from the input size
# Calculation for a 32x32 input:
# 1. Conv2d(kernel=3, stride=1, padding=0): 32x32 -> 30x30
# 2. MaxPool2d(kernel=2, stride=2): 30x30 -> 15x15
input_features = 64 * 15 * 15 # 64 is the number of output channels of conv1
self.fc = torch.nn.Linear(input_features, 10)
def forward(self, x):
x = self.conv1(x) # [B, 64, 30, 30]
x = self.relu(x)
x = self.pool(x) # [B, 64, 15, 15]
x = x.view(x.size(0), -1) # [B, 64*15*15]
return self.fc(x)
Let me first explain this model.
The Structure of SimpleModel
This code is a basic CNN model for image recognition.
The model takes a 32×32-pixel RGB image (3 channels) as input, expands the 3 input channels to 64 channels, and detects features with a 3×3 kernel.
After the convolutional layer comes a ReLU activation function, followed by a pooling layer that halves the size of the feature map by taking the maximum value from each 2×2 region.
Finally, the fully connected layer receives the features extracted by the preceding layers and classifies them into 10 classes. Before feeding into this layer, the multi-dimensional feature map is flattened into one dimension.
2.2. Converting to TorchScript
Now, before converting the CNN above to TorchScript, let's look at the conversion methods.
There are two ways to convert to TorchScript: tracing and scripting.
Although there are two methods, both are very simple to use, and the TorchScript conversion code can be written as follows
def convert_to_torchscript(model, method='trace'):
model.eval() # Set to evaluation mode
if method == 'trace':
# Tracing
example_input = torch.randn(1, 3, 32, 32)
traced_model = torch.jit.trace(model, example_input)
return traced_model
else:
# Scripting
return torch.jit.script(model)
The two approaches have the following characteristics
Tracing
- Follows the actual computation using dummy data
- Well suited to fixed input patterns
- Used when the processing flow is simple
Scripting
- Analyzes and converts the code directly
- Allows flexible logic
- Well suited to complex models with conditionals and loops
2.3. Usage Example
Now, let's convert the CNN model we built earlier into a production-ready TorchScript model.
By passing method='trace', we convert it using tracing.
# Create the model
model = SimpleModel()
# Convert for production
production_model = convert_to_torchscript(model, method='trace')
# Run inference with the converted model
input_data = torch.randn(1, 3, 32, 32)
result = production_model(input_data)
Notice that convert_to_torchscript feeds in dummy data with example_input = torch.randn(1, 3, 32, 32).
By feeding in dummy data (real production data is fine too), the network is "traced".
As you can see,
the conversion to a production-ready form is remarkably easy.
With this genuinely simple procedure, a PyTorch model created in research and development can be converted into a format that runs efficiently in production environments.
That alone gives you faster execution and a format that can run in environments that do not depend on Python.
3 TorchScript's Internal Representation
Let's peek inside a model converted with TorchScript.
I will show you how to actually inspect what is happening inside the model and what optimizations have been applied.
3.1 How to Inspect the Internal Representation
First, let's create a function to inspect the internal representation:
def inspect_torchscript(model):
# Print the model's graph structure
print("Graph Structure:")
print(model.graph)
# Print the optimized code
print("\nOptimized Code:")
print(model.code)
# Print the operators in use
print("\nOperators:")
for node in model.graph.nodes():
print(f"- {node.kind()}")
3.2 Understanding the Internal Representation
When you use this function, the first thing displayed is the model's computation flow as a graph structure. It shows how the PyTorch code we wrote is actually converted into a flow of computations. You can grasp the overall picture of how data is transformed inside the model and which layers process it in what order.
The optimized code shown next is the code after conversion into the form that is actually executed. You can see how TorchScript's optimizations changed the original code we wrote. This information is especially useful for debugging.
Finally, the list of operators tells you every kind of operation used in the model. By examining it, you can check whether any computationally heavy operators are present and find areas with room for optimization.
3.3 Practical Usage
Let's try it out.
We first create the model, convert it to TorchScript, and then look inside
model = SimpleModel()
production_model = convert_to_torchscript(model, method='script')
input_data = torch.randn(1, 3, 32, 32)
result = production_model(input_data)
print(f"Input shape: {input_data.shape}")
print(f"Output shape: {result.shape}")
inspect_torchscript(production_model)
3.4 TorchScript and the Computation Graph
The result of the run above is as follows.
Input shape: torch.Size([1, 3, 32, 32])
Output shape: torch.Size([1, 10])
Graph Structure:
graph(%self : __torch__.SimpleModel,
%x.1 : Tensor):
%17 : int = prim::Constant[value=-1]() # ex/ts1.py:21:30
%13 : int = prim::Constant[value=0]() # ex/ts1.py:21:26
%conv1 : __torch__.torch.nn.modules.conv.Conv2d = prim::GetAttr[name="conv1"](%self)
%x.5 : Tensor = prim::CallMethod[name="forward"](%conv1, %x.1) # ex/ts1.py:18:12
%relu : __torch__.torch.nn.modules.activation.ReLU = prim::GetAttr[name="relu"](%self)
%x.9 : Tensor = prim::CallMethod[name="forward"](%relu, %x.5) # ex/ts1.py:19:12
%pool : __torch__.torch.nn.modules.pooling.MaxPool2d = prim::GetAttr[name="pool"](%self)
%x.13 : Tensor = prim::CallMethod[name="forward"](%pool, %x.9) # ex/ts1.py:20:12
%14 : int = aten::size(%x.13, %13) # ex/ts1.py:21:19
%18 : int[] = prim::ListConstruct(%14, %17)
%x.19 : Tensor = aten::view(%x.13, %18) # ex/ts1.py:21:12
%fc : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="fc"](%self)
%22 : Tensor = prim::CallMethod[name="forward"](%fc, %x.19) # ex/ts1.py:22:15
return (%22)
Optimized Code:
def forward(self,
x: Tensor) -> Tensor:
conv1 = self.conv1
x0 = (conv1).forward(x, )
relu = self.relu
x1 = (relu).forward(x0, )
pool = self.pool
x2 = (pool).forward(x1, )
x3 = torch.view(x2, [torch.size(x2, 0), -1])
fc = self.fc
return (fc).forward(x3, )
Operators:
- prim::Constant
- prim::Constant
- prim::GetAttr
- prim::CallMethod
- prim::GetAttr
- prim::CallMethod
- prim::GetAttr
- prim::CallMethod
- aten::size
- prim::ListConstruct
- aten::view
- prim::GetAttr
- prim::CallMethod
As shown above, we converted SimpleModel to TorchScript and displayed its internal structure.
First, as evidence that the CNN code is implemented correctly, we confirmed that when a 32×32-pixel RGB image (shape: [1, 3, 32, 32]) is given as input, predictions for 10 classes (shape: [1, 10]) are produced as output.
Recalling the Computation Graph to Appreciate What TorchScript Does
TorchScript converts the Python code we write into a form called a computation graph.
Do you remember computation graphs?
Once you get used to PyTorch's autograd, you stop thinking about the computation graph explicitly—so let's refresh our memory a bit.
A computation graph represents the flow of operations inside a model as a directed graph. Each node represents an operation (a convolution, a ReLU, and so on), and edges represent the flow of data. This structure precisely traces the computation performed in the forward pass.
During training, we first compute the loss value in the forward pass.
Next, we need to compute the gradient of each parameter based on this loss value.
That is because we need to know how much each parameter contributes to the final loss.
Seen from the loss, the influence of the parameters in the layer immediately before it is easy to compute, but the further back a layer is, the more computations sit in between, making its direct influence harder to compute.
So by computing backwards from the loss in order, we can use the chain rule to obtain the gradient of each parameter efficiently.
That is backpropagation.
Now that it is coming back to you, let's deepen our understanding with a concrete example of the computation graph and the chain rule.
If you remember computation graphs and the chain rule perfectly well and only want to see the TorchScript results, feel free to skip the following.
A Concrete Example of a Computation Graph
Consider a computation where we take an input $x$, first multiply it by $a$, then add $b$, and finally square the result. Written as a formula:
$$
\begin{aligned}
y &= (ax + b)^2
\end{aligned}
$$
This computation can be broken down into the following small steps:
$$
\begin{aligned}
p &= ax \
q &= p + b \
y &= q^2
\end{aligned}
$$
The Chain Rule and Gradient Computation
The chain rule is what lets us decompose the gradient as follows
$$
\begin{aligned}
\frac{dy}{dx} &= \frac{dy}{dq} \times \frac{dq}{dp} \times \frac{dp}{dx} \
&= 2q \times 1 \times a \
&= 2a(ax + b)
\end{aligned}
$$
Now let's look at the gradient at each step
$$
\begin{aligned}
\frac{dy}{dq} &= 2q & \text{(derivative of squaring)} \
\frac{dq}{dp} &= 1 & \text{(derivative of addition)} \
\frac{dp}{dx} &= a & \text{(derivative of multiplication)}
\end{aligned}
$$
Differentiation tells us how the error changes when we tweak a given parameter slightly. Backpropagation is a method of computing derivatives backwards from the error, and thanks to this mechanism, no matter how complex the neural network, as long as we can compute the gradient of each individual operation (that is, differentiate it), we can efficiently obtain the overall gradient.
At the risk of belaboring the point: a computation graph can be seen as expressing this backpropagation process as a combination of simple computations.
Now, let's return to the TorchScript computation graph. In this graph, the flow of data is clearly visualized.
First, how each layer is retrieved.
Through the prim::GetAttr operation, the convolutional layer (conv1), ReLU layer (relu), pooling layer (pool), and fully connected layer (fc) are retrieved from the model in order. Each layer is then executed via the prim::CallMethod operation.
Next, the tensor reshaping operations.
aten::size obtains the batch size, prim::ListConstruct builds the new shape, and finally aten::view reshapes the tensor. This corresponds to x.view(x.size(0), -1) in the Python code.
The Optimized Code
As we have seen, TorchScript optimizes the original code and converts it into a more efficient form.
As proof, looking at the optimized code, you can see the original Python code has been converted into a more direct form.
Each layer is retrieved as a local variable, and its forward function is called.
Since this is optimized code aiming for the shortest path, Python's overhead is eliminated, enabling more efficient execution.
For example, the tensor reshaping operation has been converted to execute directly as torch.view.
The Operators Used
Finally, let's look at the operators used in this model.
In addition to basic constant creation (prim::Constant), attribute retrieval (prim::GetAttr), and method calls (prim::CallMethod), tensor operations (aten::size, aten::view) are used. These represent low-level operations, and we were able to see how the model's execution is carried out at this lower level
4 Common TorchScript Errors and How to Fix Them
TorchScript looks great, but it cannot be applied to just any code.
Write your model carelessly and it simply will not convert to TorchScript.
4.1 Errors Involving Variadic and Keyword Arguments
If you write code without thinking about it, you will see error messages like this
Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults
For example, code like the following will trigger it every time.
# Example of code that raises the error
class ProblematicModel(torch.nn.Module):
def forward(self, x, return_latents=False, **kwargs): # <- this is the error
# Processing
pass
Cause
TorchScript has the following constraints
- Variadic arguments (
*args,**kwargs) cannot be used - Keyword-only arguments with default values cannot be used
When switching between several models, it is tempting to shove parameters that only one particular model needs into forward's **kwargs. This does not play well with TorchScript, so if you want to optimize, it is safer not to write code that casually switches between multiple models whose parameters differ slightly
Fixes
As for fixes, here are a few options
- Define every argument explicitly
# Fix 1: Explicitly specify all required arguments
class FixedModel(torch.nn.Module):
def forward(self, x, return_latents=False, return_rgb=True, randomize_noise=True):
# Processing
pass
- Use a configuration class or dataclass
# Fix 2: Use a class that bundles the settings
@dataclass
class ModelConfig:
return_latents: bool = False
return_rgb: bool = True
randomize_noise: bool = True
class BetterModel(torch.nn.Module):
def forward(self, x, config: ModelConfig):
# Read the settings from config and process
pass
4.2 Implementation Best Practices
Based on the previous two sections, keep the following in mind when writing models for TorchScript
-
Define arguments explicitly
- Avoid variadic arguments (
*args,**kwargs) - If default values are needed, define them as regular arguments
- Avoid variadic arguments (
-
Use type hints
from typing import Tuple, Optional
class TypedModel(torch.nn.Module):
def forward(self, x: torch.Tensor, return_latents: bool = False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
# Processing
pass
- Use a configuration class when complex settings are needed
from dataclasses import dataclass
@dataclass
class GeneratorConfig:
return_latents: bool = False
return_rgb: bool = True
randomize_noise: bool = True
noise_scale: float = 1.0
class Generator(torch.nn.Module):
def forward(self, x: torch.Tensor, config: GeneratorConfig) -> torch.Tensor:
# Process using config
output = self.process(x, noise_scale=config.noise_scale)
if config.return_latents:
# Return latent variables
pass
return output
By applying these fixes, the conversion to TorchScript goes smoothly, and you can avoid the last-minute tragedy of discovering, right before production, that your model cannot be converted to TorchScript.
In this article, we covered PyTorch optimization—specifically how to convert Python code to TorchScript, what happens behind the scenes, a refresher on computation graphs, and TorchScript anti-patterns.
Thank you for reading to the end.
See you next time!
Appendix: The Complete Code Used in This Article
import torch
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(3, 64, 3)
self.relu = torch.nn.ReLU()
self.pool = torch.nn.MaxPool2d(2)
# Compute the input dimension of the fully connected layer from the input size
# Calculation for a 32x32 input:
# 1. Conv2d(kernel=3, stride=1, padding=0): 32x32 -> 30x30
# 2. MaxPool2d(kernel=2, stride=2): 30x30 -> 15x15
input_features = 64 * 15 * 15 # 64 is the number of output channels of conv1
self.fc = torch.nn.Linear(input_features, 10)
def forward(self, x):
x = self.conv1(x) # [B, 64, 30, 30]
x = self.relu(x)
x = self.pool(x) # [B, 64, 15, 15]
x = x.view(x.size(0), -1) # [B, 64*15*15]
return self.fc(x)
def convert_to_torchscript(model, method='trace'):
model.eval()
if method == 'trace':
example_input = torch.randn(1, 3, 32, 32)
traced_model = torch.jit.trace(model, example_input)
return traced_model
else:
return torch.jit.script(model)
def inspect_torchscript(model):
# Print the model's graph structure
print("Graph Structure:")
print(model.graph)
# Print the optimized code
print("\nOptimized Code:")
print(model.code)
# Print the operators in use
print("\nOperators:")
for node in model.graph.nodes():
print(f"- {node.kind()}")
# Run
if __name__ == "__main__":
model = SimpleModel()
production_model = convert_to_torchscript(model, method='script')
input_data = torch.randn(1, 3, 32, 32)
result = production_model(input_data)
print(f"Input shape: {input_data.shape}")
print(f"Output shape: {result.shape}")
inspect_torchscript(production_model)
