PyTorch TorchScript Tracing

Tracing records the operations from one real run of a model. The recording becomes a standalone program that repeats those operations on new inputs.

Tracing a simple model

A model with no branches traces cleanly. The traced copy doubles its input just like the original.

import torch

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

traced = torch.jit.trace(Doubler(), torch.tensor([4.0]))
print(traced(torch.tensor([4.0])).item())

The branch the example took

The example input is positive, and the recording keeps the doubling branch. A negative input still gets doubled.

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

traced = torch.jit.trace(Branchy(), torch.tensor([2.0]))
print(traced(torch.tensor([-2.0])).item())

Two inputs, one path

Both signs follow the recorded path here. The negative input doubles, where the original model would triple it.

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

traced = torch.jit.trace(Branchy(), torch.tensor([2.0]))
print(traced(torch.tensor([5.0])).item())
print(traced(torch.tensor([-5.0])).item())

Example

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

torch.manual_seed(0)
model = Branchy()
traced = torch.jit.trace(model, torch.tensor([2.0]))

print("Input | Eager | Traced | Match")
for value in [1.0, 3.0, 0.5]:
    x = torch.tensor([value])
    eager = model(x).item()
    got = traced(x).item()
    print(f"{value} | {eager:.1f} | {got:.1f} | {eager == got}")

x = torch.tensor([-2.0])
print(f"Negative input eager: {model(x).item():.1f}")
print(f"Negative input traced: {traced(x).item():.1f}")

Key points

Trace: a recording of the operations from one run of a model
Use torch.jit.trace(model, example) to build a traced model
The example input decides which branch the recording keeps
A traced model repeats the recorded path, and the path never changes
Branches on tensor values are baked in, and later inputs follow them

In this example, the Branchy model doubles positive inputs and triples negative ones. Tracing runs it once on a positive example, so the recording keeps the doubling branch. Three positive inputs then match the original model exactly. The negative input tells the truth about tracing. The original returns minus 6, and the traced copy returns minus 4. The branch decision was baked in during the recording.

Note

Torch 2.14 marks jit.trace as deprecated, and its warning points to torch.export

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

class Branchy(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * 3

torch.manual_seed(0)
model = Branchy()
traced = torch.jit.trace(model, torch.tensor([2.0]))

print("Input | Eager | Traced | Match")
for value in [1.0, 3.0, 0.5]:
    x = torch.tensor([value])
    eager = model(x).item()
    got = traced(x).item()
    print(f"{value} | {eager:.1f} | {got:.1f} | {eager == got}")

x = torch.tensor([-2.0])
print(f"Negative input eager: {model(x).item():.1f}")
print(f"Negative input traced: {traced(x).item():.1f}")

Exercise

A traced model was recorded on a positive input. What does it return for a negative input?

The recorded branch's result An error about the branch The eager model's result Always zero