ONNX Export

Sparsify, fold, export and check a model in ONNX — float, then INT8

Overview

One path, run end to end: fine-tune a ResNet-18, sparsify it, fold its batch norms, export it to ONNX and check the file that comes out. Then the same model in INT8, scored against the float file.

1. Train, then sparsify

Oxford-IIIT Pet, two classes, 64px. Three epochs of fine-tuning, then two more under SparsifyCallback with sparsity=0.5, a fraction: half the weights of every convolution end up zero.

print("torch", torch.__version__, "| onnx", onnx.__version__,
      "| onnxruntime", onnxruntime.__version__)

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64), bs=64)

torch.manual_seed(0)
net = resnet18(weights=ResNet18_Weights.DEFAULT)
net.fc = nn.Linear(512, 2)

learn = Learner(dls, net, metrics=accuracy)
with learn.no_bar(): learn.fit_one_cycle(3, 1e-3)
torch 2.9.1+cu128 | onnx 1.17.0 | onnxruntime 1.24.1
[0, 0.3779972493648529, 0.36142417788505554, 0.8552097678184509, '00:04']
[1, 0.23007234930992126, 0.19752638041973114, 0.9194858074188232, '00:04']
[2, 0.11908359825611115, 0.18409502506256104, 0.9336941838264465, '00:05']
sp_cb = SparsifyCallback(sparsity=0.5, granularity='weight', context='local',
                         criteria=large_final, schedule=one_cycle)
with learn.no_bar(): learn.fit_one_cycle(2, 1e-3, cbs=sp_cb)
Sparsifying weight until a sparsity of 50.00%
Saving Weights at epoch 0
Sparsity at the end of epoch 0: 36.57%
[0, 0.18205663561820984, 0.26363885402679443, 0.8883626461029053, '00:08']
Sparsity at the end of epoch 1: 50.00%
[1, 0.12024571001529694, 0.17602984607219696, 0.934370756149292, '00:09']
Final Sparsity: 50.00%

Sparsity Report:
--------------------------------------------------------------------------------
Layer                          Type            Params     Zeros      Sparsity  
--------------------------------------------------------------------------------
conv1                          Conv2d          9,408      4,702         49.98%
layer1.0.conv1                 Conv2d          36,864     18,430        49.99%
layer1.0.conv2                 Conv2d          36,864     18,430        49.99%
layer1.1.conv1                 Conv2d          36,864     18,430        49.99%
layer1.1.conv2                 Conv2d          36,864     18,430        49.99%
layer2.0.conv1                 Conv2d          73,728     36,862        50.00%
layer2.0.conv2                 Conv2d          147,456    73,725        50.00%
layer2.0.downsample.0          Conv2d          8,192      4,094         49.98%
layer2.1.conv1                 Conv2d          147,456    73,725        50.00%
layer2.1.conv2                 Conv2d          147,456    73,725        50.00%
layer3.0.conv1                 Conv2d          294,912    147,451       50.00%
layer3.0.conv2                 Conv2d          589,824    294,903       50.00%
layer3.0.downsample.0          Conv2d          32,768     16,382        49.99%
layer3.1.conv1                 Conv2d          589,824    294,903       50.00%
layer3.1.conv2                 Conv2d          589,824    294,903       50.00%
layer4.0.conv1                 Conv2d          1,179,648  589,808       50.00%
layer4.0.conv2                 Conv2d          2,359,296  1,179,618     50.00%
layer4.0.downsample.0          Conv2d          131,072    65,533        50.00%
layer4.1.conv1                 Conv2d          2,359,296  1,179,618     50.00%
layer4.1.conv2                 Conv2d          2,359,296  1,179,618     50.00%
--------------------------------------------------------------------------------
Overall                        all             11,166,912 5,583,290     50.00%

2. Fold batch norm

BN_Folder folds each BatchNorm2d into the convolution before it, leaving none in the model:

model = BN_Folder().fold(learn.model)
model.eval().cpu()

print(sum(1 for m in model.modules() if isinstance(m, nn.BatchNorm2d)), "BatchNorm2d modules left")
0 BatchNorm2d modules left

The batches the rest of the page reuses: four for calibration, and every full validation batch — a captured graph runs one batch size, so the short last batch is dropped.

BATCH = 64

calibration_batches = [xb.cpu() for xb, _ in islice(dls.train, 4)]
valid_batches = [(xb.cpu(), yb.cpu()) for xb, yb in dls.valid if len(yb) == BATCH]

print(len(calibration_batches), "calibration batches |", len(valid_batches), "validation batches |",
      sum(len(yb) for _, yb in valid_batches), "images")
4 calibration batches | 23 validation batches | 1472 images

3. Export the float model

export_onnx writes the file, verify_onnx compares its outputs against PyTorch’s, ONNXModel runs it through ONNX Runtime.

sample = torch.randn(1, 3, 64, 64)

float_path = export_onnx(model, sample, TMP/"model.onnx")
print(float_path.name, f"{float_path.stat().st_size / 1e6:.1f} MB")
print("verify_onnx:", verify_onnx(model, float_path, sample))
print("ONNXModel output:", tuple(ONNXModel(float_path, device="cpu")(sample).shape))
model.onnx 44.7 MB
verify_onnx: True
ONNXModel output: (1, 2)

What counts is what the file carries, not what the pipeline intended, so read the zeros back out of the exported convolution weights:

proto = onnx.load(str(float_path))
initializers = {i.name: i for i in proto.graph.initializer}
weights = [numpy_helper.to_array(initializers[n.input[1]])
           for n in proto.graph.node if n.op_type == "Conv" and n.input[1] in initializers]

zeros, total = sum(int((w == 0).sum()) for w in weights), sum(w.size for w in weights)
print(f"{len(weights)} Conv weight tensors: {zeros}/{total} zeros = {zeros/total:.2%}")
print("declared opset:", [(o.domain or "ai.onnx", o.version) for o in proto.opset_import])
20 Conv weight tensors: 5583290/11166912 zeros = 50.00%
declared opset: [('ai.onnx', 17)]

4. The same model in INT8

Quantizer(backend='pt2e') captures the model with torch.export, calibrates on those batches and converts it to a symmetric INT8 graph. export_qdq writes it with the Q/DQ pairs intact; qdq_stats counts what ended up in the file.

qmodel = Quantizer(backend='pt2e').quantize(
    model, calibration_dl=calibration_batches, max_calibration_samples=256)

spec = quant_spec(qmodel)
print(spec.as_dict())
print(spec.label, "| exportable:", spec.exports)
{'backend': 'pt2e', 'method': 'static', 'weight_bits': 8, 'act_bits': 8, 'qscheme': 'per_channel', 'symmetric': True, 'group_size': None, 'layer_bits': None, 'qdq_placement': 'per_op', 'skip_activations': None, 'engine': None}
W8A8 | exportable: True
qdq_path = export_qdq(qmodel, calibration_batches[0], TMP/"model_int8.onnx")
print(qdq_path.name, f"{qdq_path.stat().st_size / 1e6:.1f} MB",
      f"({float_path.stat().st_size / qdq_path.stat().st_size:.1f}x smaller than the float file)")
print(qdq_stats(qdq_path).as_dict())
print("declared opset:", [(o.domain or "ai.onnx", o.version) for o in onnx.load(str(qdq_path)).opset_import])
[torch.onnx] Obtain model graph for `GraphModule([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `GraphModule([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
model_int8.onnx 11.4 MB (3.9x smaller than the float file)
{'n_quantize': 33, 'n_dequantize': 54, 'n_per_channel': 21, 'n_nonzero_zero_point': 0, 'n_unquantized_conv_add': 0}
declared opset: [('ai.onnx', 18)]

The two files declare different opsets because those are the two library defaults: export_onnx asks for opset 17 and export_qdq for opset 18, the opset the Q/DQ translations it registers are written against.

def report(name, k, n, z=1.96):
    "A count with its n and its Wilson 95% interval"
    p, d, centre = k/n, 1 + z**2/n, k/n + z**2/(2*n)
    half = z*((p*(1-p) + z**2/(4*n))/n)**0.5
    print(f"{name}: {k}/{n} = {p:.4f}  Wilson 95% [{(centre-half)/d:.4f}, {(centre+half)/d:.4f}]")

def score(onnx_path, batches):
    "Correct predictions and n for an ONNX file, over the batches it is given"
    onnx_model = ONNXModel(onnx_path, device="cpu")
    correct = sum(int((onnx_model(xb).argmax(-1) == yb).sum()) for xb, yb in batches)
    return correct, sum(len(yb) for _, yb in batches)

def score_torch(torch_model, batches):
    "The same count for a torch model, over the same batches"
    torch_model.eval()
    with torch.no_grad():
        correct = sum(int((torch_model(xb).argmax(-1) == yb).sum()) for xb, yb in batches)
    return correct, sum(len(yb) for _, yb in batches)

verify_qdq reports how often the exported graph and the quantized model agree on the class. That is a parity guard only if the reference predictions vary, so count them first:

probe = torch.cat([xb for xb, _ in valid_batches[:4]])  # 256 real validation images

with torch.no_grad():
    reference = torch.cat([qmodel(xb).argmax(-1) for xb, _ in valid_batches[:4]])
print("reference predictions per class:", torch.bincount(reference, minlength=2).tolist())

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter('always')
    agreement = verify_qdq(qmodel, qdq_path, probe, n_batches=4)
print("vacuity warning:", any('vacuous' in str(w.message) for w in caught))
report("argmax agreement", round(agreement * len(probe)), len(probe))
reference predictions per class: [165, 91]
vacuity warning: False
argmax agreement: 256/256 = 1.0000  Wilson 95% [0.9852, 1.0000]

The exported graph and the quantized model agree on every probe input, and the reference predicts both classes rather than one: a parity guard, not the agreement a single-class reference returns whatever the graph computes.

The two files and the PyTorch model they came from, scored on the same batches. verify_onnx above compares the file’s output against PyTorch’s on one random batch-1 input within a tolerance, which is not an accuracy, so the torch model is scored the same way as the files rather than read off that check:

report("FP32 PyTorch", *score_torch(model, valid_batches))
report("FP32 ONNX", *score(float_path, valid_batches))
report("INT8 QDQ ONNX", *score(qdq_path, valid_batches))
FP32 PyTorch: 1375/1472 = 0.9341  Wilson 95% [0.9203, 0.9457]
FP32 ONNX: 1375/1472 = 0.9341  Wilson 95% [0.9203, 0.9457]
INT8 QDQ ONNX: 1373/1472 = 0.9327  Wilson 95% [0.9188, 0.9444]

Three rows over the same batches: the model in PyTorch, the float file, and the INT8 file. PyTorch and the float file come out on the same count, which is what the export is asked to preserve; the INT8 row’s Wilson interval overlaps both, so this single run does not order the three.

Measured on

ResNet-18 (ImageNet weights, fc replaced), Oxford-IIIT Pet (two classes), 64px, batch 64; three epochs of fine-tuning, two under SparsifyCallback, batch norms folded. torch 2.9.1+cu128 / onnx 1.17.0 / onnxruntime 1.24.1, CPU, single run, one seed. The three accuracies cover the same 23 full validation batches, n = 1472: the PyTorch row runs the folded, sparsified model in eager mode, the other two go through ONNX Runtime. The parity check uses the first four batches, n = 256. No latency is measured, and no non-sparse baseline is scored.

Summary

Tool What it gives you
SparsifyCallback(sparsity=0.5, ...) Half of every convolution’s weights zeroed during the fit, with a per-layer report
BN_Folder().fold(model) The same model with every BatchNorm2d folded into its convolution
export_onnx(model, sample, path) A float ONNX file, dynamic batch, opset 17
verify_onnx(model, path, sample) Whether the file’s outputs match PyTorch’s within tolerance
ONNXModel(path) A callable that runs the file through ONNX Runtime
Quantizer(backend='pt2e') A symmetric INT8 graph, calibrated on the batches you hand it
export_qdq(model, sample, path) An ONNX file whose Q/DQ pairs survived the export, opset 18
qdq_stats(path) Node counts, per-channel count and zero-point audit, read off the file
verify_qdq(model, path, samples) Argmax agreement between the exported graph and PyTorch

Every compression ratio on this page is a fraction in [0, 1]: sparsity=0.5 is half the weights.


See Also