Just-in-time (JIT) compilation translates or optimizes model code during execution. In PyTorch, it keeps eager-mode development convenient while compiling models into graphs for improved production performance. section 17.3.1
PyTorch uses several graph-compilation approaches:
- torch.jit.trace records operations performed on sample inputs, producing TorchScript, but struggles with input-dependent control flow. section 17.3.1
- torch.jit.script analyzes more Python-like loops and conditionals, but supports only a subset of Python and can fail at runtime. section 17.3.1
- torch.compile uses TorchDynamo to capture dynamic graphs, splitting graphs and falling back when it encounters unrecognized control paths. section 17.3.1
A model with dynamic control flow
A model can choose between two linear layers according to whether x.sum() is positive. Exporting this model reports that dynamic control flow is unsupported, whereas torch.compile(model) provides a compilation path while preserving the eager-style model definition. section 17.3.1
Listing 17.10 A model with dynamic conditional control flow
import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class ConditionalModel(nn.Module):
def __init__(self):
super(ConditionalModel, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(10, 3)
def forward(self, x):
if x.sum() > 0:
x = self.fc1(x)
else:
x = self.fc2(x)
x = self.relu(x)
return x
model = ConditionalModel()
model = model.to(device)
model.eval()
Inspect a complete ConditionalModel that selects layers through an input-dependent if statement before applying ReLU. section 17.3.1
