Chapter 4: PyTorch TorchScript

Production Deployment for PyTorch Models

弘益人間 · Benefit All Humanity

4.1 Introduction to TorchScript

TorchScript is PyTorch's way to create serializable and optimizable models from PyTorch code. It bridges the gap between research (Python) and production (C++), enabling deployment in environments without Python dependencies. TorchScript uses JIT (Just-In-Time) compilation to optimize models for inference.

The core philosophy of TorchScript is to provide a path from dynamic, eager execution in Python to static, optimized execution in production environments. This transition is crucial for deploying models at scale, where performance, reliability, and minimal dependencies are paramount.

4.2 Tracing vs Scripting

TorchScript offers two methods to convert PyTorch models: tracing and scripting. Understanding when to use each method is critical for successful model deployment.

Tracing

Tracing records operations executed during a forward pass with example inputs. It's ideal for models with static computation graphs.

import torch

# Define model
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        return self.relu(self.fc(x))

model = SimpleModel()
example_input = torch.randn(1, 10)

# Trace the model
traced_model = torch.jit.trace(model, example_input)

# Save
traced_model.save('traced_model.pt')

# Load and use
loaded = torch.jit.load('traced_model.pt')
output = loaded(torch.randn(1, 10))

# Inspect the traced graph
print(traced_model.graph)
print(traced_model.code)

Scripting

Scripting analyzes the Python source code directly, capturing control flow and dynamic behavior. Use this when your model contains if statements, loops, or other control structures.

class ConditionalModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)

    def forward(self, x):
        # Control flow that must be preserved
        if x.sum() > 0:
            return self.fc(x)
        else:
            return x[:, :5]

model = ConditionalModel()

# Script the model (captures control flow)
scripted_model = torch.jit.script(model)

# Save
scripted_model.save('scripted_model.pt')

# Verify control flow is preserved
print(scripted_model.code)

4.3 When to Use Tracing vs Scripting

Choose the right method based on your model's characteristics:

4.4 Advanced TorchScript Features

Type Annotations

Type annotations are essential for scripting, helping TorchScript understand your code's intent and catch errors early.

import torch
from typing import List, Tuple, Dict, Optional

class AnnotatedModel(torch.nn.Module):
    def forward(
        self,
        x: torch.Tensor,
        hidden: Tuple[torch.Tensor, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        h, c = hidden
        output = self.process(x, h)
        return {"output": output, "hidden": c}

    def process(self, x: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
        return x + h

    def batch_process(
        self,
        inputs: List[torch.Tensor],
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Process a batch with optional masking"""
        result = torch.stack(inputs)
        if mask is not None:
            result = result * mask
        return result

Custom Operators

import torch

@torch.jit.script
def custom_activation(x: torch.Tensor, threshold: float) -> torch.Tensor:
    """Custom gated activation function"""
    return torch.where(x > threshold, x, torch.zeros_like(x))

@torch.jit.script
def swish(x: torch.Tensor, beta: float = 1.0) -> torch.Tensor:
    """Swish activation: x * sigmoid(beta * x)"""
    return x * torch.sigmoid(beta * x)

class ModelWithCustomOp(torch.nn.Module):
    def forward(self, x):
        x = custom_activation(x, 0.5)
        x = swish(x, beta=1.5)
        return x

model = ModelWithCustomOp()
scripted = torch.jit.script(model)

4.5 Loading and Using TorchScript Models

In Python

import torch
import numpy as np

# Load model
model = torch.jit.load('model.pt')
model.eval()

# Run inference
with torch.no_grad():
    output = model(input_tensor)

# Inspect the graph
print("Graph structure:")
print(model.graph)

print("\nGenerated code:")
print(model.code)

# Check model parameters
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

# Benchmark inference
import time
start = time.time()
for _ in range(100):
    with torch.no_grad():
        _ = model(input_tensor)
elapsed = (time.time() - start) / 100
print(f"Average inference time: {elapsed*1000:.2f}ms")

In C++

#include <torch/script.h>
#include <iostream>
#include <memory>
#include <chrono>

int main() {
    // Load model
    torch::jit::script::Module module;
    try {
        module = torch::jit::load("model.pt");
        module.eval();
    } catch (const c10::Error& e) {
        std::cerr << "Error loading model: " << e.what() << "\n";
        return -1;
    }

    // Create input
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::randn({1, 3, 224, 224}));

    // Execute model with timing
    auto start = std::chrono::high_resolution_clock::now();
    at::Tensor output = module.forward(inputs).toTensor();
    auto end = std::chrono::high_resolution_clock::now();

    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
    std::cout << "Inference time: " << duration.count() << "ms\n";
    std::cout << "Output shape: " << output.sizes() << '\n';
    std::cout << "Output (first 5): " << output.slice(1, 0, 5) << '\n';

    return 0;
}

// Compile with:
// g++ -std=c++17 inference.cpp -o inference \
//     -I${TORCH_PATH}/include \
//     -L${TORCH_PATH}/lib \
//     -ltorch -ltorch_cpu -lc10

4.6 Optimization Techniques

Graph Optimization

import torch

model = torch.jit.load('model.pt')

# Optimize for inference
optimized_model = torch.jit.optimize_for_inference(model)

# Freeze model (inline constants, remove training-only ops)
frozen_model = torch.jit.freeze(optimized_model)

# Additional optimizations
# 1. Remove dropout layers
frozen_model.eval()

# 2. Fuse operations
# Conv + BatchNorm + ReLU are automatically fused

# Save optimized version
frozen_model.save('optimized_model.pt')

# Compare model sizes
import os
original_size = os.path.getsize('model.pt') / (1024 * 1024)
optimized_size = os.path.getsize('optimized_model.pt') / (1024 * 1024)
print(f"Original: {original_size:.2f}MB")
print(f"Optimized: {optimized_size:.2f}MB")
print(f"Reduction: {(1 - optimized_size/original_size)*100:.1f}%")

Operator Fusion

PyTorch automatically fuses common operator patterns to reduce memory bandwidth and improve performance.

# PyTorch automatically fuses operations
# Conv + BatchNorm + ReLU → Fused operation

model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 64, 3),
    torch.nn.BatchNorm2d(64),
    torch.nn.ReLU()
)

traced = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
# Fusion happens automatically during tracing

# Common fusion patterns:
# 1. Conv + BN + ReLU
# 2. Linear + ReLU
# 3. Add + ReLU
# 4. Mul + Add (fused multiply-add)

# View fused graph
print(traced.graph)

4.7 Mobile Deployment

PyTorch Mobile

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torch.jit.load('model.pt')
model.eval()

# Optimize for mobile
mobile_model = optimize_for_mobile(model)

# Save for lite interpreter (smaller runtime)
mobile_model._save_for_lite_interpreter('mobile_model.ptl')

# Verify mobile model
loaded_mobile = torch.jit.load('mobile_model.ptl')
test_input = torch.randn(1, 3, 224, 224)
output = loaded_mobile(test_input)
print(f"Mobile model output shape: {output.shape}")

# In Android (Java):
"""
import org.pytorch.LiteModuleLoader;
import org.pytorch.Module;
import org.pytorch.Tensor;

Module module = LiteModuleLoader.load("mobile_model.ptl");
Tensor input = Tensor.fromBlob(inputArray, shape);
Tensor output = module.forward(IValue.from(input)).toTensor();
float[] scores = output.getDataAsFloatArray();
"""

# In iOS (Swift):
"""
import LibTorch

guard let module = try? LiteModuleLoader.load(modelPath: "mobile_model.ptl") else {
    fatalError("Failed to load model")
}
let input = try! Tensor(shape: [1, 3, 224, 224])
let output = try! module.forward([input])
"""

4.8 Quantization

Quantization reduces model size and improves inference speed by converting weights and activations from floating point to integers.

import torch

# Post-training static quantization
model.eval()

# Step 1: Fuse modules
model_fused = torch.ao.quantization.fuse_modules(
    model,
    [['conv', 'bn', 'relu']]
)

# Step 2: Specify quantization config
model_fused.qconfig = torch.ao.quantization.get_default_qconfig('x86')

# For mobile deployment:
# model_fused.qconfig = torch.ao.quantization.get_default_qconfig('qnnpack')

# Step 3: Prepare for quantization
model_prepared = torch.ao.quantization.prepare(model_fused)

# Step 4: Calibrate with sample data
with torch.no_grad():
    for data, _ in calibration_loader:
        model_prepared(data)

# Step 5: Convert to quantized model
quantized_model = torch.ao.quantization.convert(model_prepared)

# Save quantized model
torch.jit.save(torch.jit.script(quantized_model), 'quantized.pt')

# Compare sizes and performance
original_size = sum(p.numel() * p.element_size() for p in model.parameters())
quantized_size = sum(p.numel() * p.element_size() for p in quantized_model.parameters())
print(f"Size reduction: {original_size / quantized_size:.1f}x")

# Benchmark
import time
inputs = torch.randn(10, 3, 224, 224)

start = time.time()
with torch.no_grad():
    _ = model(inputs)
fp32_time = time.time() - start

start = time.time()
with torch.no_grad():
    _ = quantized_model(inputs)
int8_time = time.time() - start

print(f"Speed improvement: {fp32_time / int8_time:.1f}x")

4.9 Best Practices

Handle Dynamic Shapes

class DynamicModel(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        seq_len = x.size(1)
        # Use dynamic dimensions appropriately
        return self.process(x, batch_size, seq_len)

    def process(self, x: torch.Tensor, bs: int, seq: int) -> torch.Tensor:
        # Avoid hard-coded shapes
        output = torch.zeros(bs, seq, 512, device=x.device)
        return output + x

# Trace with different input sizes to verify
traced = torch.jit.trace(model, torch.randn(1, 10, 512))

# Test with various batch sizes
assert traced(torch.randn(4, 20, 512)).shape[0] == 4
assert traced(torch.randn(8, 15, 512)).shape[0] == 8

# Use torch.jit.script for models with truly dynamic behavior
scripted = torch.jit.script(model)

Error Handling

try:
    model = torch.jit.script(MyModel())
    print("Scripting successful!")
except RuntimeError as e:
    print(f"Scripting failed: {e}")
    print("Attempting tracing instead...")
    try:
        example_input = torch.randn(1, 3, 224, 224)
        model = torch.jit.trace(MyModel(), example_input)
        print("Tracing successful!")
    except Exception as e:
        print(f"Both scripting and tracing failed: {e}")
        # Consider fixing model code or using eager mode

# Validate converted model
def validate_model(original, converted, test_inputs):
    """Validate that converted model matches original"""
    original.eval()
    converted.eval()

    with torch.no_grad():
        for inp in test_inputs:
            orig_out = original(inp)
            conv_out = converted(inp)
            diff = torch.abs(orig_out - conv_out)
            max_diff = torch.max(diff).item()
            print(f"Max difference: {max_diff:.2e}")
            assert max_diff < 1e-5, f"Outputs differ too much: {max_diff}"

test_inputs = [torch.randn(1, 3, 224, 224) for _ in range(5)]
validate_model(original_model, traced_model, test_inputs)

4.10 Debugging TorchScript Models

Debugging TorchScript models requires different techniques than standard Python debugging.

Inspect Generated Code

import torch

# Load or create scripted model
model = torch.jit.script(MyModel())

# View the generated TorchScript code
print(model.code)

# View the computation graph
print(model.graph)

# Get detailed graph representation
for node in model.graph.nodes():
    print(f"Node: {node.kind()}")
    print(f"  Inputs: {[i.debugName() for i in node.inputs()]}")
    print(f"  Outputs: {[o.debugName() for o in node.outputs()]}")

# Check for specific operations
def find_operations(graph, op_name):
    """Find all nodes of a specific operation type"""
    return [node for node in graph.nodes() if node.kind() == op_name]

conv_ops = find_operations(model.graph, 'aten::conv2d')
print(f"Found {len(conv_ops)} convolution operations")

Performance Profiling

import torch
from torch.profiler import profile, ProfilerActivity

model = torch.jit.load('model.pt')
input_data = torch.randn(1, 3, 224, 224)

# Profile with PyTorch profiler
with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    profile_memory=True
) as prof:
    with torch.no_grad():
        model(input_data)

# Print profiling results
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

# Export Chrome trace for visualization
prof.export_chrome_trace("trace.json")
# Open trace.json in chrome://tracing