06: Exporting Models for Deployment¶

Every previous tutorial ends with a trained (and possibly pruned/quantized) nn.Module sitting inside a live Python process. That is fine for experimentation, but most real deployment targets do not run a Python interpreter at all: a mobile app, a browser, an embedded device, a low-latency C++ inference server, or a serving stack like ONNX Runtime/TensorRT/CoreML.

Export is the step that turns a dynamic PyTorch model into a static, Python-free artifact those runtimes can load. shrinkai.export provides functions for that.

This notebook will show:

When to reach for ONNX vs. TorchScript.
Export to ONNX.
Export to TorchScript.
How export interacts with everything shrinkai.compression produces.
Export directly from a Distiller.

Setup¶

Copied!





import contextlib
import copy
import ctypes
import os

import torch
import torch.nn as nn

from shrinkai.compression.pruning import ChannelPruner, Pruner, PruningConfig
from shrinkai.compression.quantization import QuantConfig, Quantizer
from shrinkai.distillation import Distiller
from shrinkai.export import export_onnx, export_torchscript

_libc = ctypes.CDLL(None)


@contextlib.contextmanager
def quiet_onnx_export_diagnostics():
    """Silences a PyTorch internal diagnostic that dumps the *entire* traced
    IR graph to stdout when `torch.onnx.export` fails. It writes through a
    buffered C stream, so a plain `contextlib.redirect_stdout` does not catch
    it — this redirects the real file descriptor and flushes libc's buffer
    before restoring it, so the dump lands in `/dev/null` instead of here.
    """
    saved_fd = os.dup(1)
    devnull_fd = os.open(os.devnull, os.O_WRONLY)
    os.dup2(devnull_fd, 1)
    try:
        yield
    finally:
        _libc.fflush(None)
        os.dup2(saved_fd, 1)
        os.close(devnull_fd)
        os.close(saved_fd)
import contextlib
import copy
import ctypes
import os

import torch
import torch.nn as nn

from shrinkai.compression.pruning import ChannelPruner, Pruner, PruningConfig
from shrinkai.compression.quantization import QuantConfig, Quantizer
from shrinkai.distillation import Distiller
from shrinkai.export import export_onnx, export_torchscript

_libc = ctypes.CDLL(None)


@contextlib.contextmanager
def quiet_onnx_export_diagnostics():
    """Silences a PyTorch internal diagnostic that dumps the *entire* traced
    IR graph to stdout when `torch.onnx.export` fails. It writes through a
    buffered C stream, so a plain `contextlib.redirect_stdout` does not catch
    it — this redirects the real file descriptor and flushes libc's buffer
    before restoring it, so the dump lands in `/dev/null` instead of here.
    """
    saved_fd = os.dup(1)
    devnull_fd = os.open(os.devnull, os.O_WRONLY)
    os.dup2(devnull_fd, 1)
    try:
        yield
    finally:
        _libc.fflush(None)
        os.dup2(saved_fd, 1)
        os.close(devnull_fd)
        os.close(saved_fd)

Copied!





def build_teacher_student():
    teacher = torch.hub.load(
        "chenyaofo/pytorch-cifar-models", "cifar10_vgg19_bn", pretrained=True, verbose=False
    )
    student = torch.hub.load(
        "chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=True, verbose=False
    )
    return teacher.eval(), student.eval()
def build_teacher_student():
    teacher = torch.hub.load(
        "chenyaofo/pytorch-cifar-models", "cifar10_vgg19_bn", pretrained=True, verbose=False
    )
    student = torch.hub.load(
        "chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=True, verbose=False
    )
    return teacher.eval(), student.eval()

Unlike previous tutorials, we do not need a training loop here: export only cares about the model's structure, not its accuracy. We reuse the pretrained VGG19-BN/ResNet20 pair from Tutorials 01/04/05 — both already trained — and stand in student for "whatever distiller.fit(...) just produced".

Copied!

teacher, student = build_teacher_student()
sample_input = torch.randn(1, 3, 32, 32)

print(f"Student: {sum(p.numel() for p in student.parameters()):,} parameters")
teacher, student = build_teacher_student()
sample_input = torch.randn(1, 3, 32, 32)

print(f"Student: {sum(p.numel() for p in student.parameters()):,} parameters")

Student: 272,474 parameters

Both exporting formats "freeze" the model by recording the sequence of tensor operations it performs, either by tracing a real forward pass, or (for method="script", see below) by compiling its Python source through a restricted subset-of-Python compiler. What differs is where the result can run afterwards:

	ONNX (`export_onnx`)	TorchScript (`export_torchscript`)
Serialization	Protobuf graph, standardized operator set ("opset")	PyTorch's own IR (`.pt`)
Runtimes	ONNX Runtime, TensorRT, CoreML/NNAPI converters, and any other stack that speaks ONNX	LibTorch (C++), PyTorch Mobile
Cross-framework	Yes — designed as a framework-agnostic exchange format	No — PyTorch-specific, by PyTorch, for PyTorch
Quantized models (`shrinkai.compression.quantization`)	Not supported by this exporter, see below	Supported natively

Rule of thumb: reach for ONNX when the target runtime is not PyTorch at all (ONNX Runtime, TensorRT, a mobile conversion pipeline), and on the other hand, reach for TorchScript when the target is LibTorch/PyTorch Mobile, or when the model uses quantized ops that ONNX's opset does not cover.

Exporting to ONNX¶

The function export_onnx(model, sample_input, path) runs sample_input through the model once to record its graph, then writes it to a .onnx file. It requires the optional onnx package.

pip install shrinkai[export]

Copied!





import onnx

onnx_path = export_onnx(student, sample_input, "exported/student.onnx")
onnx_model = onnx.load(str(onnx_path))
onnx.checker.check_model(onnx_model)  # structural validation, raises if the graph is malformed

print(f"Exported to {onnx_path} ({onnx_path.stat().st_size / 1024:.1f} KB)")
print(f"Graph has {len(onnx_model.graph.node)} nodes, {len(onnx_model.graph.initializer)} initializers")
import onnx

onnx_path = export_onnx(student, sample_input, "exported/student.onnx")
onnx_model = onnx.load(str(onnx_path))
onnx.checker.check_model(onnx_model)  # structural validation, raises if the graph is malformed

print(f"Exported to {onnx_path} ({onnx_path.stat().st_size / 1024:.1f} KB)")
print(f"Graph has {len(onnx_model.graph.node)} nodes, {len(onnx_model.graph.initializer)} initializers")

Exported to exported/student.onnx (1071.7 KB)
Graph has 59 nodes, 44 initializers

Note: We do not have onnxruntime installed here, so we validate the export structurally (onnx.checker + shape inference below) rather than by re-running the model and comparing outputs numerically. For a true end-to-end numeric check, install onnxruntime and call onnxruntime.InferenceSession(str(onnx_path)).run(...).

By default (dynamic_batch=True), export_onnx marks dimension 0 of every input as dynamic, so the exported graph accepts any batch size at inference time, not just the one sample_input happened to have. ONNX's own shape inference then propagates that symbolic dimension to the outputs.

Copied!





from onnx import shape_inference

dynamic_path = export_onnx(student, sample_input, "exported/student_dynamic.onnx", input_names=["input"])
static_path = export_onnx(
    student, sample_input, "exported/student_static.onnx", input_names=["input"], dynamic_batch=False
)

dynamic_model = shape_inference.infer_shapes(onnx.load(str(dynamic_path)))
static_model = onnx.load(str(static_path))

dynamic_dim0 = dynamic_model.graph.output[0].type.tensor_type.shape.dim[0]
static_dim0 = static_model.graph.input[0].type.tensor_type.shape.dim[0]

print("dynamic_batch=True  -> output batch dim:", dynamic_dim0.dim_param or dynamic_dim0.dim_value, "(symbolic)")
print("dynamic_batch=False -> input batch dim:", static_dim0.dim_value, "(frozen to sample_input's batch size)")
from onnx import shape_inference

dynamic_path = export_onnx(student, sample_input, "exported/student_dynamic.onnx", input_names=["input"])
static_path = export_onnx(
    student, sample_input, "exported/student_static.onnx", input_names=["input"], dynamic_batch=False
)

dynamic_model = shape_inference.infer_shapes(onnx.load(str(dynamic_path)))
static_model = onnx.load(str(static_path))

dynamic_dim0 = dynamic_model.graph.output[0].type.tensor_type.shape.dim[0]
static_dim0 = static_model.graph.input[0].type.tensor_type.shape.dim[0]

print("dynamic_batch=True  -> output batch dim:", dynamic_dim0.dim_param or dynamic_dim0.dim_value, "(symbolic)")
print("dynamic_batch=False -> input batch dim:", static_dim0.dim_value, "(frozen to sample_input's batch size)")

dynamic_batch=True  -> output batch dim: batch_size (symbolic)
dynamic_batch=False -> input batch dim: 1 (frozen to sample_input's batch size)

Exporting to TorchScript¶

export_torchscript supports two distinct compilation strategies, and the choice matters:

method="trace" (default) runs sample_input through the model once and records the exact tensor operations that were executed fast, and works on almost any model, but it "bakes in" whichever code path ran for that specific input. Any data-dependent Python control flow (an if on a tensor's value, a loop whose length depends on the input) is silently frozen to whatever branch was taken during tracing.
method="script" compiles the model's actual Python source through TorchScript's subset-of-Python compiler. Slower to get working (the source must be expressible in that subset) and more likely to reject unsupported constructs, but it preserves control flow exactly.

In this example, since student (ResNet20) has no data-dependent branching, both methods work and agree with each other and with the eager model.

Copied!





traced_path = export_torchscript(student, "exported/student_traced.pt", sample_input=sample_input, method="trace")
traced = torch.jit.load(str(traced_path))

with torch.no_grad():
    eager_out = student(sample_input)
    traced_out = traced(sample_input)

print(f"traced .pt size: {traced_path.stat().st_size / 1024:.1f} KB")
print("traced matches eager:", torch.allclose(eager_out, traced_out, atol=1e-5))
traced_path = export_torchscript(student, "exported/student_traced.pt", sample_input=sample_input, method="trace")
traced = torch.jit.load(str(traced_path))

with torch.no_grad():
    eager_out = student(sample_input)
    traced_out = traced(sample_input)

print(f"traced .pt size: {traced_path.stat().st_size / 1024:.1f} KB")
print("traced matches eager:", torch.allclose(eager_out, traced_out, atol=1e-5))

traced .pt size: 1240.8 KB
traced matches eager: True

Copied!





scripted_path = export_torchscript(student, "exported/student_scripted.pt", method="script")
scripted = torch.jit.load(str(scripted_path))

with torch.no_grad():
    scripted_out = scripted(sample_input)

print(f"scripted .pt size: {scripted_path.stat().st_size / 1024:.1f} KB")
print("scripted matches eager:", torch.allclose(eager_out, scripted_out, atol=1e-5))
scripted_path = export_torchscript(student, "exported/student_scripted.pt", method="script")
scripted = torch.jit.load(str(scripted_path))

with torch.no_grad():
    scripted_out = scripted(sample_input)

print(f"scripted .pt size: {scripted_path.stat().st_size / 1024:.1f} KB")
print("scripted matches eager:", torch.allclose(eager_out, scripted_out, atol=1e-5))

scripted .pt size: 1150.7 KB
scripted matches eager: True

The difference above is invisible on ResNet20 because it has no data-dependent branching. Here is a minimal model that does, to make the trap concrete: a forward that picks between two Linear layers based on the sign of the input's sum, a genuine data-dependent branch.

Copied!





class ControlFlowModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.positive_branch = nn.Linear(4, 4)
        self.negative_branch = nn.Linear(4, 4)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.sum() > 0:
            return self.positive_branch(x)
        else:
            return self.negative_branch(x)


toy_model = ControlFlowModel().eval()
positive_input = torch.ones(1, 4)
negative_input = -torch.ones(1, 4)

# Traced while `positive_input` takes the `if` branch: that branch gets baked in permanently.
toy_traced_path = export_torchscript(
    toy_model, "exported/toy_traced.pt", sample_input=positive_input, method="trace"
)
toy_scripted_path = export_torchscript(toy_model, "exported/toy_scripted.pt", method="script")

toy_traced = torch.jit.load(str(toy_traced_path))
toy_scripted = torch.jit.load(str(toy_scripted_path))

with torch.no_grad():
    print("eager,   negative input:", toy_model(negative_input))
    print("traced,  negative input:", toy_traced(negative_input), " <- wrong branch, baked in at trace time")
    print("scripted, negative input:", toy_scripted(negative_input), " <- correct, control flow preserved")
class ControlFlowModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.positive_branch = nn.Linear(4, 4)
        self.negative_branch = nn.Linear(4, 4)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.sum() > 0:
            return self.positive_branch(x)
        else:
            return self.negative_branch(x)


toy_model = ControlFlowModel().eval()
positive_input = torch.ones(1, 4)
negative_input = -torch.ones(1, 4)

# Traced while `positive_input` takes the `if` branch: that branch gets baked in permanently.
toy_traced_path = export_torchscript(
    toy_model, "exported/toy_traced.pt", sample_input=positive_input, method="trace"
)
toy_scripted_path = export_torchscript(toy_model, "exported/toy_scripted.pt", method="script")

toy_traced = torch.jit.load(str(toy_traced_path))
toy_scripted = torch.jit.load(str(toy_scripted_path))

with torch.no_grad():
    print("eager,   negative input:", toy_model(negative_input))
    print("traced,  negative input:", toy_traced(negative_input), " <- wrong branch, baked in at trace time")
    print("scripted, negative input:", toy_scripted(negative_input), " <- correct, control flow preserved")

eager,   negative input: tensor([[ 0.6323,  0.3464,  0.2897, -1.0566]])
traced,  negative input: tensor([[ 0.0651, -0.1629,  0.8471, -0.0015]])  <- wrong branch, baked in at trace time
scripted, negative input: tensor([[ 0.6323,  0.3464,  0.2897, -1.0566]])  <- correct, control flow preserved

/var/folders/sf/rj62rfhs7mg6zs95s_xxg25r0000gn/T/ipykernel_34236/1593689705.py:8: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if x.sum() > 0:

The traced export silently returns the wrong answer on negative_input: it was traced on positive_input, so the if branch got recorded as an unconditional call to positive_branch, no matter what the input actually is later. PyTorch even warns about this at trace time, it is worth never silencing that warning away. However, method="script" has no such issue, since it compiles the actual if statement instead of recording one execution of it.

To keep in mind: trace is the right default for models without data-dependent control flow (the vast majority of CNNs/Transformers doing straight-line computation). Reach for script when the model branches on tensor values or has input-dependent loop lengths.

Exporting compressed models¶

Export composes with everything in shrinkai.compression, but each compression technique interacts with it a little differently.

Mask-based Pruner¶

Pruner.apply() attaches torch.nn.utils.prune forward hooks that multiply the weight by a mask on every forward pass. The unpruned weight and the mask are both still stored as separate tensors. Exporting through those hooks works, but it traces the multiply itself into the graph and ships both tensors, doubling the relevant parameter storage for nothing. Pruner.finalize() bakes the mask into the weight and removes the hooks: always call it before export.

Copied!

pruned = Pruner(PruningConfig(method="unstructured", amount=0.4)).apply(copy.deepcopy(student))

unfinalized_path = export_onnx(pruned, sample_input, "exported/pruned_unfinalized.onnx")
finalized_path = export_onnx(Pruner.finalize(pruned), sample_input, "exported/pruned_finalized.onnx")

unfinalized_model = onnx.load(str(unfinalized_path))
finalized_model = onnx.load(str(finalized_path))

print(f"before finalize(): {len(unfinalized_model.graph.node)} nodes, {unfinalized_path.stat().st_size / 1024:.1f} KB")
print(f"after  finalize(): {len(finalized_model.graph.node)} nodes, {finalized_path.stat().st_size / 1024:.1f} KB")
pruned = Pruner(PruningConfig(method="unstructured", amount=0.4)).apply(copy.deepcopy(student))

unfinalized_path = export_onnx(pruned, sample_input, "exported/pruned_unfinalized.onnx")
finalized_path = export_onnx(Pruner.finalize(pruned), sample_input, "exported/pruned_finalized.onnx")

unfinalized_model = onnx.load(str(unfinalized_path))
finalized_model = onnx.load(str(finalized_path))

print(f"before finalize(): {len(unfinalized_model.graph.node)} nodes, {unfinalized_path.stat().st_size / 1024:.1f} KB")
print(f"after  finalize(): {len(finalized_model.graph.node)} nodes, {finalized_path.stat().st_size / 1024:.1f} KB")

before finalize(): 82 nodes, 1089.0 KB
after  finalize(): 59 nodes, 1071.7 KB

Fewer nodes and a smaller file, for the exact same (masked) predictions, purely because the export no longer carries the mask-multiply and the duplicate weight tensor.

ChannelPruner¶

ChannelPruner physically rebuilds smaller Conv2d/Linear/BatchNorm layers, thus, the result is an ordinary nn.Module with no hooks and no masks, so it exports exactly like any other model:

Copied!





channel_pruned = ChannelPruner(amount=0.3).apply(
    copy.deepcopy(student),
    target_layers=["layer1.0.conv1", "layer2.0.conv1"],
    sample_input=sample_input,
)

original_path = export_onnx(student, sample_input, "exported/original.onnx")
channel_pruned_path = export_onnx(channel_pruned, sample_input, "exported/channel_pruned.onnx")

print(f"original:       {original_path.stat().st_size / 1024:.1f} KB")
print(f"channel-pruned: {channel_pruned_path.stat().st_size / 1024:.1f} KB")
channel_pruned = ChannelPruner(amount=0.3).apply(
    copy.deepcopy(student),
    target_layers=["layer1.0.conv1", "layer2.0.conv1"],
    sample_input=sample_input,
)

original_path = export_onnx(student, sample_input, "exported/original.onnx")
channel_pruned_path = export_onnx(channel_pruned, sample_input, "exported/channel_pruned.onnx")

print(f"original:       {original_path.stat().st_size / 1024:.1f} KB")
print(f"channel-pruned: {channel_pruned_path.stat().st_size / 1024:.1f} KB")

original:       1071.7 KB
channel-pruned: 1051.9 KB

Quantizer¶

Eager-mode quantized layers (torch.ao.nn.quantized.*) compile to TorchScript natively, that is precisely the deployment path PyTorch designed them for. PyTorch's legacy ONNX exporter, however, does not know, how to translate those quantized ops (quantized::linear_dynamic,quantized::batch_norm2d, ...) into ONNX's operator set. This is a real, current limitation, not a shrinkai bug, we show it happening rather than paper over it.

Note: These two cells quantize teacher (VGG19-BN), not student (ResNet20). As Tutorial 04 covers in depth, fully converting a residual architecture end-to-end hits a separate, unrelated limitation of PyTorch's eager-mode quantization (out += identity has no QuantizedCPU kernel), nothing to do with export. VGG has no residual connections, so it isolates the export behavior we actually want to demonstrate here.

Copied!





torch.backends.quantized.engine = "qnnpack"

dynamic_quantized = Quantizer(QuantConfig(strategy="ptq")).apply(copy.deepcopy(teacher))
dynamic_quantized.eval()

# TorchScript: works.
dq_trace_path = export_torchscript(
    dynamic_quantized, "exported/dynamic_quantized.pt", sample_input=sample_input, method="trace"
)
dq_traced = torch.jit.load(str(dq_trace_path))
with torch.no_grad():
    dq_matches = torch.allclose(dynamic_quantized(sample_input), dq_traced(sample_input))
print(f"dynamic PTQ -> TorchScript: OK, {dq_trace_path.stat().st_size / 1024:.1f} KB, matches eager: {dq_matches}")

# ONNX: fails, and we surface exactly why (the huge IR graph PyTorch dumps
# alongside the exception is suppressed by quiet_onnx_export_diagnostics —
# only the actual error message matters here).
try:
    with quiet_onnx_export_diagnostics():
        export_onnx(dynamic_quantized, sample_input, "exported/dynamic_quantized.onnx")
except Exception as exc:
    print(f"dynamic PTQ -> ONNX: {type(exc).__name__}: {str(exc).strip().splitlines()[-1]}")
torch.backends.quantized.engine = "qnnpack"

dynamic_quantized = Quantizer(QuantConfig(strategy="ptq")).apply(copy.deepcopy(teacher))
dynamic_quantized.eval()

# TorchScript: works.
dq_trace_path = export_torchscript(
    dynamic_quantized, "exported/dynamic_quantized.pt", sample_input=sample_input, method="trace"
)
dq_traced = torch.jit.load(str(dq_trace_path))
with torch.no_grad():
    dq_matches = torch.allclose(dynamic_quantized(sample_input), dq_traced(sample_input))
print(f"dynamic PTQ -> TorchScript: OK, {dq_trace_path.stat().st_size / 1024:.1f} KB, matches eager: {dq_matches}")

# ONNX: fails, and we surface exactly why (the huge IR graph PyTorch dumps
# alongside the exception is suppressed by quiet_onnx_export_diagnostics —
# only the actual error message matters here).
try:
    with quiet_onnx_export_diagnostics():
        export_onnx(dynamic_quantized, sample_input, "exported/dynamic_quantized.onnx")
except Exception as exc:
    print(f"dynamic PTQ -> ONNX: {type(exc).__name__}: {str(exc).strip().splitlines()[-1]}")

dynamic PTQ -> TorchScript: OK, 78994.6 KB, matches eager: True
dynamic PTQ -> ONNX: UnsupportedOperatorError: Exporting the operator 'quantized::linear_dynamic' to ONNX opset version 17 is not supported

Note the exported dynamic-PTQ TorchScript file is not much smaller than the original: Quantizer's dynamic PTQ path only quantizes nn.Linear layers (see Quantizer._apply_ptq), and VGG19-BN is almost entirely convolutions, which is the exact same observation made in Tutorial 04. Dynamic PTQ shines on Linear-heavy models; for conv-heavy ones, Static PTQ or ChannelPruner are the better fit. Static PTQ hits the same ONNX wall, one op earlier.

Copied!





static_quantized = Quantizer(QuantConfig(strategy="ptq", calibrate_data=[sample_input, sample_input])).apply(
    copy.deepcopy(teacher)
)
static_quantized.eval()

sq_trace_path = export_torchscript(
    static_quantized, "exported/static_quantized.pt", sample_input=sample_input, method="trace"
)
sq_traced = torch.jit.load(str(sq_trace_path))
with torch.no_grad():
    sq_matches = torch.allclose(static_quantized(sample_input), sq_traced(sample_input))
print(f"static PTQ -> TorchScript: OK, {sq_trace_path.stat().st_size / 1024:.1f} KB, matches eager: {sq_matches}")

try:
    with quiet_onnx_export_diagnostics():
        export_onnx(static_quantized, sample_input, "exported/static_quantized.onnx")
except Exception as exc:
    print(f"static PTQ -> ONNX: {type(exc).__name__}: {str(exc).strip().splitlines()[-1]}")
static_quantized = Quantizer(QuantConfig(strategy="ptq", calibrate_data=[sample_input, sample_input])).apply(
    copy.deepcopy(teacher)
)
static_quantized.eval()

sq_trace_path = export_torchscript(
    static_quantized, "exported/static_quantized.pt", sample_input=sample_input, method="trace"
)
sq_traced = torch.jit.load(str(sq_trace_path))
with torch.no_grad():
    sq_matches = torch.allclose(static_quantized(sample_input), sq_traced(sample_input))
print(f"static PTQ -> TorchScript: OK, {sq_trace_path.stat().st_size / 1024:.1f} KB, matches eager: {sq_matches}")

try:
    with quiet_onnx_export_diagnostics():
        export_onnx(static_quantized, sample_input, "exported/static_quantized.onnx")
except Exception as exc:
    print(f"static PTQ -> ONNX: {type(exc).__name__}: {str(exc).strip().splitlines()[-1]}")

static PTQ -> TorchScript: OK, 20384.9 KB, matches eager: True
static PTQ -> ONNX: UnsupportedOperatorError: Exporting the operator 'quantized::batch_norm2d' to ONNX opset version 17 is not supported

To keep in mind: if a model must be both quantized and deployed on a non-PyTorch runtime, this legacy ONNX exporter is not the path, but deploy it via TorchScript/LibTorch instead, or look into ONNX Runtime's own native quantization tooling as a separate route that does not go through PyTorch's quantized ops at all.

Export directly from a Distiller¶

Everything above is also reachable directly from a Distiller instance. Indeed, distiller.export_onnx(...)/distiller.export_torchscript(...) simply forward to shrinkai.export.export_onnx/export_torchscript on distiller.student, so there is nothing new to learn, only fewer imports:

Copied!

distiller = Distiller(teacher=teacher, student=student, device="cpu")
# distiller.fit(train_loader, val_loader, epochs=...)  # skipped here — see Tutorials 01/05

onnx_export_path = distiller.export_onnx("exported/from_distiller.onnx", sample_input)
ts_export_path = distiller.export_torchscript("exported/from_distiller.pt", sample_input)

print("via Distiller ->", onnx_export_path, "and", ts_export_path)
distiller = Distiller(teacher=teacher, student=student, device="cpu")
# distiller.fit(train_loader, val_loader, epochs=...)  # skipped here — see Tutorials 01/05

onnx_export_path = distiller.export_onnx("exported/from_distiller.onnx", sample_input)
ts_export_path = distiller.export_torchscript("exported/from_distiller.pt", sample_input)

print("via Distiller ->", onnx_export_path, "and", ts_export_path)

via Distiller -> exported/from_distiller.onnx and exported/from_distiller.pt

Summary¶

Need	Reach for
Deploy on ONNX Runtime, TensorRT, or a mobile conversion pipeline	`export_onnx`
Deploy via LibTorch (C++) or PyTorch Mobile	`export_torchscript(method="trace")`
Preserve data-dependent control flow exactly	`export_torchscript(method="script")`
Accept any batch size at inference time	`export_onnx(..., dynamic_batch=True)` (default)
Export a mask-pruned model (`Pruner`)	`Pruner.finalize()` first, then export
Export a channel-pruned model (`ChannelPruner`)	export directly, no extra step
Export a quantized model (`Quantizer`)	`export_torchscript` only — ONNX unsupported here
One-liner export from a `Distiller`	`distiller.export_onnx` / `distiller.export_torchscript`

References¶

Bai, J., Lu, F., Zhang, K., et al. (2019). ONNX: Open Neural Network Exchange. Linux Foundation.
https://github.com/onnx/onnx
PyTorch documentation (2024). torch.onnx: export internals, supported operators, and opset versions.
https://pytorch.org/docs/stable/onnx.html
PyTorch documentation (2024). TorchScript: torch.jit.trace vs. torch.jit.script semantics.
https://pytorch.org/docs/stable/jit.html
PyTorch documentation (2024). Quantization: which ops are TorchScript-compatible after quantization.
https://pytorch.org/docs/stable/quantization.html