01 · Exporting Models 📦

Your model was born inside a Python training loop. To deploy it, you often need to set it free — into a portable artifact that runs without your code, or even without Python at all.

---

1. The 'Why'

Everything you've built so far assumes a friendly environment: a Python interpreter, your model class definition importable, PyTorch installed, maybe a GPU. Production is rarely so accommodating. You might need to run your model inside a C++ application, on a mobile phone, in a browser, on an edge device with no Python, or on a high-throughput server where spinning up the Python interpreter for every request is too slow. In all these cases you can't just ship your .py files and a state_dict — you need to package the model into a self-contained, portable artifact that captures both the computation and the weights, and that other runtimes can execute directly. That packaging step is exporting.

The reason exporting is its own skill (and not just "save the file") is that a PyTorch model in eager mode is really Python code — a forward method full of arbitrary Python that runs line by line. To run it elsewhere, you must capture that dynamic computation as a static, language-independent graph. PyTorch offers a few ways to do this, each a different trade-off between flexibility and portability: TorchScript (the long-standing approach, via tracing or scripting), the newer torch.export (PyTorch 2.x's graph capture), and ONNX (an open, framework-neutral format supported by many runtimes). You don't need mastery of all three today — you need to understand what exporting is, why it's necessary, and which tool to reach for.

---

2. Core Concepts

Eager mode vs. a captured graph. Normally PyTorch runs your forward eagerly, one Python op at a time. Exporting captures that computation into a graph that can run without the Python source.

TorchScript: tracing vs. scripting. torch.jit.trace runs an example input through the model and records the ops — simple, but it misses data-dependent control flow (if/loops that vary by input). torch.jit.script compiles the code itself, preserving control flow. Traced/scripted modules save to a single .pt file runnable from Python or C++ (LibTorch).

torch.export. The modern PyTorch 2.x path: produces a clean, standardized graph (an ExportedProgram) intended as the foundation for downstream deployment and optimization. It's stricter but more principled than tracing.

ONNX. An open, framework-agnostic model format. Exporting to ONNX (torch.onnx.export) lets your model run in ONNX Runtime and many other engines/hardware backends, independent of PyTorch.

Export in eval() mode with a representative input. Always switch to eval() first (so Dropout/BatchNorm behave for inference), and provide an example input whose shape matches production.

---

3. Code in Action

TorchScript via tracing

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
model.eval()                                  # ALWAYS export in eval mode

example = torch.randn(1, 4)                    # a representative input shape
traced = torch.jit.trace(model, example)       # record the ops for this input
traced.save("model_traced.pt")                 # self-contained; runs without the class

# Reload anywhere (Python or C++/LibTorch) — no original class needed:
loaded = torch.jit.load("model_traced.pt")
print(loaded(torch.randn(1, 4)).shape)         # torch.Size([1, 3])

TorchScript via scripting (preserves control flow)

import torch
import torch.nn as nn

class Conditional(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(4, 3)
    def forward(self, x):
        if x.sum() > 0:            # data-dependent branch — tracing would miss this
            return self.fc(x)
        return self.fc(x) * 0.5

scripted = torch.jit.script(Conditional().eval())  # compiles the code, branches intact
scripted.save("model_scripted.pt")

Exporting to ONNX

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3)).eval()
example = torch.randn(1, 4)

torch.onnx.export(
    model, example, "model.onnx",
    input_names=["input"], output_names=["logits"],
    dynamic_axes={"input": {0: "batch"}},      # allow a variable batch dimension
)
# model.onnx now runs in ONNX Runtime and many other engines — no PyTorch required.

---

4. Common Pitfalls

Pitfall 1 — Tracing a model with data-dependent control flow. torch.jit.trace only records the path your example input took, silently baking in one branch of any if/loop that depends on the data. If your forward has such control flow, use torch.jit.script (or torch.export) instead, or the exported model will be subtly wrong.

Pitfall 2 — Exporting in training mode. Forgetting model.eval() before export bakes training-time Dropout/BatchNorm behavior into the artifact, so the deployed model gives noisy, wrong predictions. Always eval() first — the same discipline from Module 06's inference recipe.

Pitfall 3 — Hard-coding the input shape and forgetting dynamic axes. A model traced/exported with a fixed batch size may reject other batch sizes at serving time. Declare dynamic_axes (ONNX) or use dynamic-shape options so the artifact accepts the input shapes production actually sends.

---

5. Further Reading & Watch List

• 📖 Deep Learning with PyTorch, Ch. 15 — deploying to production
• 🌐 Learn PyTorch — Saving & exporting (workflow)
• 🛠️ TorchScript intro · torch.export · Exporting to ONNX

---

➡️ Next: 02 · torch.compile