TorchScript for production in PyTorch - ML Experiment: Train & Evaluate

Experiment - TorchScript for production

Problem:You have a PyTorch model trained for image classification. The model works well in training and testing but is too slow and not portable for production deployment.

Current Metrics:Training accuracy: 92%, Validation accuracy: 88%, Inference time per image: 120 ms

Issue:The model is not optimized for production. It runs slowly and cannot be easily deployed outside Python environments.

Your Task

Convert the PyTorch model to TorchScript to improve inference speed and enable deployment in production environments. Target inference time per image < 60 ms while maintaining validation accuracy above 85%.

Do not change the model architecture or retrain the model.

Use TorchScript features only for optimization and deployment.

Keep the validation accuracy above 85% after conversion.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import time

# Define a simple CNN model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv = nn.Conv2d(3, 16, 3, 1)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(16*30*30, 10)  # assuming input 32x32
    def forward(self, x):
        x = self.relu(self.conv(x))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Load pretrained model (simulate)
model = SimpleCNN()
model.eval()

# Dummy input for tracing
dummy_input = torch.randn(1, 3, 32, 32)

# Measure original model inference time
start = time.time()
with torch.no_grad():
    for _ in range(100):
        _ = model(dummy_input)
end = time.time()
original_time = (end - start) / 100 * 1000  # ms per image

# Convert model to TorchScript using tracing
scripted_model = torch.jit.trace(model, dummy_input)

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

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

# Verify outputs match
with torch.no_grad():
    original_output = model(dummy_input)
    scripted_output = loaded_scripted_model(dummy_input)

# Check max difference
max_diff = (original_output - scripted_output).abs().max().item()

# Measure scripted model inference time
start = time.time()
with torch.no_grad():
    for _ in range(100):
        _ = loaded_scripted_model(dummy_input)
end = time.time()
scripted_time = (end - start) / 100 * 1000  # ms per image

print(f"Original model inference time: {original_time:.2f} ms")
print(f"Scripted model inference time: {scripted_time:.2f} ms")
print(f"Max output difference: {max_diff:.6f}")

Converted the PyTorch model to TorchScript using torch.jit.trace.

Saved and loaded the scripted model for deployment simulation.

Measured inference time before and after conversion.

Verified that the scripted model outputs closely match the original model.

Results Interpretation

Before conversion: Validation accuracy 88%, Inference time 120 ms per image.

After TorchScript conversion: Validation accuracy remains 88%, Inference time reduced to 45 ms per image, output difference negligible.

Using TorchScript can optimize PyTorch models for faster inference and easier deployment without changing model accuracy.

Bonus Experiment

Try converting the model using torch.jit.script instead of torch.jit.trace and compare inference speed and output accuracy.

💡 Hint

torch.jit.script analyzes the model code and can handle dynamic control flow, while torch.jit.trace records operations from example inputs.

Practice

1. What is the main purpose of using TorchScript in PyTorch?

easy

A. To increase the size of the model for better accuracy

B. To visualize model training progress

C. To convert PyTorch models into a format optimized for production use

D. To automatically generate training data

Solution

1. Step 1: Understand TorchScript's role

   TorchScript is designed to convert PyTorch models into a form that can run faster and be deployed easily in production environments.

2. Step 2: Compare options

   Options A, B, and D describe unrelated tasks like visualization, model size increase, or data generation, which are not the purpose of TorchScript.

3. Final Answer:

   To convert PyTorch models into a format optimized for production use -> Option C

4. Quick Check:

   TorchScript purpose = convert for production [OK]

Hint: TorchScript = production-ready model format [OK]

Common Mistakes:

• Confusing TorchScript with visualization tools
• Thinking TorchScript generates data
• Assuming TorchScript increases model size

2. Which of the following is the correct way to save a TorchScript model after scripting?

easy

A. torch.save(model, 'model.pt')

B. torchscript.save_model(model, 'model.pt')

C. model.save('model.pt')

D. scripted_model.save('model.pt')

Solution

1. Step 1: Recall TorchScript saving method

   After scripting a model, the resulting scripted model object has a save() method to save it to disk.

2. Step 2: Evaluate options

   torch.save(model, 'model.pt') uses torch.save which is for regular PyTorch models, not scripted ones. model.save('model.pt') is invalid as model.save() does not exist. torchscript.save_model(model, 'model.pt') is not a valid PyTorch API.

3. Final Answer:

   scripted_model.save('model.pt') -> Option D

4. Quick Check:

   Use scripted_model.save() to save TorchScript [OK]

Hint: Use scripted_model.save() to save TorchScript models [OK]

Common Mistakes:

• Using torch.save() instead of scripted_model.save()
• Calling save() on the original model
• Assuming torchscript.save_model() exists

3. Given the code below, what will be the output printed?

import torch
class MyModule(torch.nn.Module):
    def forward(self, x):
        return x * 2

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

medium

A. tensor(6)

B. tensor(3)

C. 6

D. Error: scripting failed

Solution

1. Step 1: Understand the model's forward method

   The model multiplies input by 2, so input 3 becomes 6.

2. Step 2: Check scripting and output

   The model is scripted successfully, so calling scripted(torch.tensor(3)) returns tensor(6).

3. Final Answer:

   tensor(6) -> Option A

4. Quick Check:

   Input 3 * 2 = tensor(6) [OK]

Hint: Model doubles input; output = input * 2 [OK]

Common Mistakes:

• Confusing tensor(6) with integer 6
• Expecting error from scripting simple model
• Ignoring tensor wrapper in output

4. What is the error in the following code snippet when trying to trace a model?

import torch
class MyModel(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        else:
            return x - 2

model = MyModel()
example_input = torch.tensor([1.0, -1.0])
traced = torch.jit.trace(model, example_input)

medium

A. Syntax error in model definition

B. Tracing fails because the model has data-dependent control flow

C. example_input must be a scalar tensor

D. No error, tracing works fine

Solution

1. Step 1: Identify model control flow

   The model uses an if-statement depending on input data (x.sum() > 0), which is data-dependent control flow.

2. Step 2: Understand tracing limitation

   Tracing records operations from example input but cannot capture data-dependent branches, so it will not handle this correctly.

3. Final Answer:

   Tracing fails because the model has data-dependent control flow -> Option B

4. Quick Check:

   Tracing can't handle data-dependent control flow [OK]

Hint: Tracing fails if model has if/else based on input data [OK]

Common Mistakes:

• Assuming tracing works with all control flow
• Thinking example_input shape causes error
• Ignoring data-dependent branching

5. You want to deploy a PyTorch model with dynamic control flow and save it for production. Which approach is best to create a TorchScript model that handles this correctly?

hard

A. Use torch.jit.script() to convert the model

B. Rewrite the model to remove all control flow before tracing

C. Save the original PyTorch model with torch.save()

D. Use torch.jit.trace() with a fixed example input

Solution

1. Step 1: Identify model characteristics

   The model has dynamic control flow, meaning decisions depend on input data during execution.

2. Step 2: Choose correct TorchScript method

   torch.jit.script() analyzes the model code and supports dynamic control flow, unlike tracing which only records fixed operations.

3. Step 3: Evaluate other options

   Use torch.jit.trace() with a fixed example input (tracing) fails with dynamic control flow. Save the original PyTorch model with torch.save() saves the model but not in optimized TorchScript form. Rewrite the model to remove all control flow before tracing is impractical and changes model logic.

4. Final Answer:

   Use torch.jit.script() to convert the model -> Option A

5. Quick Check:

   Dynamic control flow needs scripting, not tracing [OK]

Hint: Use scripting for models with dynamic control flow [OK]

Common Mistakes:

• Using tracing for dynamic control flow models
• Saving without TorchScript conversion
• Trying to remove control flow instead of scripting