TorchScript for production in PyTorch

Introduction

TorchScript helps you turn PyTorch models into a form that runs fast and works well in real apps.

You want to run your PyTorch model on a server without Python.

You need faster model execution for real-time predictions.

You want to share your model with others without sharing Python code.

You want to deploy your model on mobile devices or embedded systems.

You want to save your model in a format that is easy to load and run anywhere.

Syntax

PyTorch

import torch

# Convert a PyTorch model to TorchScript
scripted_model = torch.jit.script(your_model)

# Save the scripted model
scripted_model.save('model.pt')

# Load the scripted model
loaded_model = torch.jit.load('model.pt')

# Use the loaded model for prediction
output = loaded_model(input_tensor)

torch.jit.script converts your model to TorchScript by analyzing the code.

You can save and load the scripted model easily for production use.

Examples

This example shows a simple model doubled input, then scripted and run.

PyTorch

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted = torch.jit.script(model)
print(scripted(torch.tensor(3)))

Tracing records operations from example input, useful for simple models.

PyTorch

import torch

# Using tracing instead of scripting
model = torch.nn.ReLU()
example_input = torch.randn(1)
traced_model = torch.jit.trace(model, example_input)
print(traced_model(torch.tensor([-1.0])))

Sample Model

This program creates a simple model, converts it to TorchScript, saves and loads it, then runs a prediction.

PyTorch

import torch
import torch.nn as nn

# Define a simple model
class MyModel(nn.Module):
    def forward(self, x):
        return x * 3 + 1

# Create model instance
model = MyModel()

# Convert to TorchScript
scripted_model = torch.jit.script(model)

# Save the scripted model
scripted_model.save('mymodel.pt')

# Load the scripted model
loaded_model = torch.jit.load('mymodel.pt')

# Test input
input_tensor = torch.tensor(4)

# Get output from loaded model
output = loaded_model(input_tensor)
print(f'Output: {output.item()}')

Important Notes

TorchScript models run faster and can be used without Python.

Use scripting for models with control flow; use tracing for simple, straight-line models.

Always test the scripted model to ensure it behaves like the original.

Summary

TorchScript converts PyTorch models for fast, production-ready use.

You can save and load TorchScript models easily.

Use scripting or tracing depending on your model's complexity.