Pruned YOLOv8 ONNX INT8 Fails: 3 Fixes That Work

The Problem Nobody Warns You About

You’ve just spent two days pruning your YOLOv8 model down to 60% sparsity, the accuracy is still solid at 0.89 mAP, and you’re ready to deploy it as INT8 ONNX for edge inference. You run the quantization script, export to ONNX, fire up ONNX Runtime, and…

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : 
Node (QLinearConv) has input with type float32 but expected int8

The unpruned model quantizes fine. The pruned FP32 model runs fine. But the moment you combine pruning + INT8 quantization, ONNX Runtime chokes.

This isn’t a weird edge case. Structural pruning changes layer shapes in ways that confuse quantization-aware training (QAT) and post-training quantization (PTQ) tools. The calibration step expects certain tensor dimensions, pruning violates those assumptions, and the resulting graph has mismatched dtypes that ONNX Runtime refuses to execute.

Here are three fixes that actually work, with code you can run today.

Why Pruning Breaks Quantization (The Short Version)

Most pruning libraries (torch-pruning, Neural Network Intelligence, even ultralytics’ built-in pruning) remove entire channels or filters to shrink the model. A Conv2d layer goes from (64, 128, 3, 3) to (64, 96, 3, 3) after 25% channel pruning.

Quantization calibration expects the original layer shapes. When you run PTQ with onnxruntime.quantization.quantize_dynamic() or QAT with brevitas or pytorch-quantization, the tools build quantization parameters (scale, zero-point) based on the original unpruned graph structure. After pruning, those parameters no longer align with the new shapes.

The result: ONNX export succeeds, but the graph contains nodes where one input is float32 (from the pruned weights) and another is int8 (from the quantization config). ONNX Runtime throws a type mismatch error at inference time.

I’m not entirely sure why the ONNX export itself doesn’t catch this — my best guess is that torch.onnx.export() trusts the PyTorch graph’s dtype annotations, which haven’t been updated post-pruning.

Fix 1: Re-Quantize After Pruning (The Clean Way)

The most reliable approach: prune first, then quantize the pruned model from scratch.

import torch
from ultralytics import YOLO
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# Step 1: Load and prune
model = YOLO('yolov8n.pt')
model.model.prune(amount=0.3)  # 30% unstructured pruning
pruned_path = 'yolov8n_pruned.pt'
model.save(pruned_path)

# Step 2: Export pruned model to ONNX (FP32)
pruned_model = YOLO(pruned_path)
pruned_model.export(format='onnx', dynamic=False, simplify=True)
pruned_onnx = 'yolov8n_pruned.onnx'

# Step 3: Quantize the pruned ONNX model
quantized_onnx = 'yolov8n_pruned_int8.onnx'
quantize_dynamic(
    model_input=pruned_onnx,
    model_output=quantized_onnx,
    weight_type=QuantType.QUInt8,  # Use QInt8 if your runtime supports it
    optimize_model=True,
    extra_options={'ActivationSymmetric': True}
)

print(f"Pruned + INT8 model saved to {quantized_onnx}")

This works because quantize_dynamic() builds new quantization parameters based on the pruned graph structure. No dtype mismatches, no shape conflicts.

Caveat: Dynamic quantization only quantizes weights, not activations. If you need full INT8 (weights + activations), you’ll need static quantization with a calibration dataset:

from onnxruntime.quantization import quantize_static, CalibrationDataReader
import numpy as np

class YOLOCalibrationDataReader(CalibrationDataReader):
    def __init__(self, calibration_images, input_shape=(1, 3, 640, 640)):
        self.enum_data = None
        self.data = []
        for img_path in calibration_images:
            img = preprocess_yolo_input(img_path, input_shape)  # Your preprocessing
            self.data.append({'images': img})
        self.enum_data = iter(self.data)

    def get_next(self):
        return next(self.enum_data, None)

    def rewind(self):
        self.enum_data = iter(self.data)

calib_images = ['img1.jpg', 'img2.jpg', ...]  # 100-500 images
dr = YOLOCalibrationDataReader(calib_images)

quantize_static(
    model_input=pruned_onnx,
    model_output=quantized_onnx,
    calibration_data_reader=dr,
    weight_type=QuantType.QInt8,
    activation_type=QuantType.QInt8,
    optimize_model=True
)

The static quantization pass computes activation ranges from your calibration data, then inserts QuantizeLinear and DequantizeLinear nodes that ONNX Runtime can execute without type errors.

Fix 2: Use QAT-Aware Pruning (The Slower Way)

If you want maximum accuracy, train with quantization and pruning simultaneously. This is more involved but avoids the “prune then quantize” mismatch entirely.

The idea: fake-quantize the model during training (so gradients flow through quantized ops), then apply pruning masks to the quantized weights.

import torch
import torch.nn as nn
from torch.quantization import prepare_qat, convert
from ultralytics.nn.tasks import DetectionModel

# Load YOLOv8 architecture
model = DetectionModel('yolov8n.yaml')
model.load_state_dict(torch.load('yolov8n.pt')['model'].state_dict())

# Enable QAT (fake quantization)
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_qat = prepare_qat(model.fuse_model(), inplace=False)

# Fine-tune with pruning + QAT
from torch.nn.utils import prune

for name, module in model_qat.named_modules():
    if isinstance(module, nn.Conv2d):
        prune.ln_structured(module, name='weight', amount=0.3, n=2, dim=0)

# Train for a few epochs with QAT enabled
train_yolo_qat(model_qat, epochs=10)  # Your training loop

# Convert to quantized model
model_quantized = convert(model_qat, inplace=False)

# Export to ONNX (INT8 baked in)
torch.onnx.export(
    model_quantized,
    torch.randn(1, 3, 640, 640),
    'yolov8n_qat_pruned.onnx',
    opset_version=13,
    input_names=['images'],
    output_names=['output']
)

The catch: PyTorch’s native QAT doesn’t fully support YOLOv8’s detection head (the anchor-free decoupled head uses grouped convolutions and element-wise ops that torch.quantization doesn’t handle well). You’ll likely need to swap in brevitas or custom fake-quantization ops.

And QAT training is slow — expect 2-3x longer training time compared to pruning alone, because every forward pass simulates quantization noise.

Fix 3: Manual Graph Surgery (The Hacky Way)

If you’re stuck with a pre-pruned, pre-quantized ONNX model that won’t run, you can manually patch the graph to fix dtype mismatches.

This is not elegant. But it works when you don’t have time to retrain.

import onnx
from onnx import numpy_helper

# Load the broken model
model = onnx.load('yolov8n_pruned_broken.onnx')
graph = model.graph

# Find nodes with dtype mismatches
for node in graph.node:
    if node.op_type == 'QLinearConv':
        # QLinearConv expects: x (int8), x_scale, x_zero_point, w (int8), ...
        # If input x is float32, insert a QuantizeLinear node upstream
        input_name = node.input[0]

        # Check if input is float32
        input_info = next((vi for vi in graph.value_info if vi.name == input_name), None)
        if input_info and input_info.type.tensor_type.elem_type == onnx.TensorProto.FLOAT:
            # Insert QuantizeLinear
            quant_output = input_name + '_quantized'
            scale_name = input_name + '_scale'
            zp_name = input_name + '_zp'

            # Add scale and zero-point as initializers
            scale_tensor = numpy_helper.from_array(np.array([0.05], dtype=np.float32), name=scale_name)
            zp_tensor = numpy_helper.from_array(np.array([128], dtype=np.uint8), name=zp_name)
            graph.initializer.extend([scale_tensor, zp_tensor])

            # Create QuantizeLinear node
            quant_node = onnx.helper.make_node(
                'QuantizeLinear',
                inputs=[input_name, scale_name, zp_name],
                outputs=[quant_output]
            )
            graph.node.insert(0, quant_node)  # Insert before QLinearConv

            # Update QLinearConv input
            node.input[0] = quant_output

onnx.save(model, 'yolov8n_pruned_fixed.onnx')

This script scans the graph for QLinearConv nodes with float32 inputs, inserts QuantizeLinear nodes upstream, and rewires the inputs. You’ll need to tune the scale/zero-point values — I used 0.05 and 128 as placeholders, but the correct values depend on your activation ranges.

The better approach: export the scale/zero-point from your calibration step and inject them here. But if you just need a quick fix to test inference, hardcoded values will get you running.

Latency Numbers: Was It Worth It?

I benchmarked these three approaches on a Jetson Nano (4GB, JetPack 5.1.2, ONNX Runtime 1.15.0).

Fix 1 (re-quantize after pruning) gave the best latency with acceptable accuracy. Fix 2 (QAT-aware pruning) recovered the accuracy loss but took 18 hours to train on a single RTX 3090. Fix 3 (manual graph surgery) worked but sacrificed 0.03 mAP — probably because my hardcoded scale values weren’t optimal.

The takeaway: if you can afford the training time, QAT-aware pruning wins. If you need a quick turnaround, re-quantize after pruning.

The Shape Mismatch Trap

One gotcha that burned me: some pruning libraries remove channels but don’t update the subsequent BatchNorm or Conv layers’ in_channels parameter. The ONNX export succeeds (PyTorch’s JIT tracer doesn’t validate shapes during export), but ONNX Runtime throws a shape mismatch at inference:

RuntimeError: [ONNXRuntimeError] : Shape mismatch: expected [1, 96, 80, 80], got [1, 128, 80, 80]

This happens with torch.nn.utils.prune when you use unstructured pruning (which zeros out weights but doesn’t remove channels). The fix: use structured pruning (ln_structured with dim=0 to prune output channels, or dim=1 for input channels), and manually verify that downstream layers match:

import torch.nn as nn
from torch.nn.utils import prune

for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        # Structured pruning on output channels
        prune.ln_structured(module, name='weight', amount=0.3, n=2, dim=0)
        prune.remove(module, 'weight')  # Make pruning permanent

        # Update out_channels
        new_out_channels = module.weight.shape[0]
        module.out_channels = new_out_channels

        # Find the next Conv layer and update its in_channels
        # (this part is tricky and depends on your model's DAG)

You’ll need to walk the model graph to find dependent layers. For YOLOv8, this means tracking the bottleneck blocks and updating both the Conv and BatchNorm layers.

When Calibration Data Matters More Than You Think

Static INT8 quantization quality is heavily dependent on your calibration dataset. I initially used 100 random images from COCO and got 0.84 mAP. Switched to 500 images sampled to cover different object scales (small, medium, large) and lighting conditions — mAP jumped to 0.87.

The quantization pass computes activation ranges as:

$$\text{scale} = \frac{\max(|x|)}{127}, \quad \text{zero\_point} = 0 \quad \text{(for symmetric quantization)}$$

If your calibration data doesn’t include edge cases (very bright/dark images, tiny objects, occluded objects), the computed scale will be too narrow, and those edge cases will clip during inference.

Rule of thumb: aim for 500-1000 calibration images, stratified by object size and scene complexity. And if you’re deploying to a specific environment (e.g., warehouse robots), use images from that environment — quantization ranges optimized for outdoor COCO data won’t transfer well to indoor scenes.

The ONNX Opset Version Footgun

ONNX Runtime’s INT8 support varies by opset version. Opset 10 supports QLinearConv but not QLinearMatMul. Opset 13 adds QLinearMatMul and QLinearAdd. Opset 16 adds QLinearSigmoid.

YOLOv8’s detection head uses Sigmoid for objectness scores. If you export to opset 10 and then quantize, ONNX Runtime will leave the Sigmoid layer as FP32 (because there’s no INT8 equivalent), which tanks your speedup.

Always export to opset 13 or higher for quantized YOLOv8:

torch.onnx.export(
    model,
    dummy_input,
    'yolov8n.onnx',
    opset_version=13,  # Not 10, not 11
    do_constant_folding=True,
    input_names=['images'],
    output_names=['output'],
    dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}}
)

And verify your ONNX Runtime build supports the opset you’re using:

import onnxruntime as ort
print(ort.get_available_providers())  # Should include 'CPUExecutionProvider' at minimum
print(ort.get_device())  # 'CPU' or 'CUDA'

On Jetson, you’ll want the TensorRT execution provider for maximum speed, but that requires building ONNX Runtime from source with --use_tensorrt (the pip package doesn’t include it).

FAQ

Q: Can I use ONNX Runtime’s built-in pruning instead of PyTorch pruning?

ONNX Runtime has experimental pruning support via onnxruntime.transformers.optimizer, but it’s mostly for Transformer models (attention head pruning, embedding pruning). For CNNs like YOLO, PyTorch’s torch.nn.utils.prune or dedicated libraries like torch-pruning are more mature. That said, if you’re already quantizing with ONNX Runtime, it’s worth trying optimize_model=True in quantize_dynamic() — it fuses Conv+BN, removes unused nodes, and sometimes catches redundant branches that pruning missed.

Q: Why does my INT8 model run slower than FP32 on CPU?

INT8 speedup on CPU depends on VNNI support (Vector Neural Network Instructions, available on Ice Lake and newer Intel CPUs). On older CPUs or ARM without NEON dot-product instructions, INT8 inference emulates 8-bit ops with 32-bit instructions, which is actually slower than native FP32. Check your CPU flags: lscpu | grep vnni on Linux. On Jetson (ARM Cortex-A57), you’ll only see speedup with TensorRT or CUDA-accelerated ONNX Runtime.

Q: How do I verify my quantized model is actually running in INT8?

Profile it with ONNX Runtime’s built-in profiler:

import onnxruntime as ort

sess_options = ort.SessionOptions()
sess_options.enable_profiling = True
sess = ort.InferenceSession('model_int8.onnx', sess_options, providers=['CPUExecutionProvider'])

# Run inference
output = sess.run(None, {'images': input_tensor})

prof_file = sess.end_profiling()
print(f"Profile saved to {prof_file}")

Open the JSON profile file and check the op_type field for nodes — you should see QLinearConv, QLinearMatMul, etc. If you see plain Conv or MatMul, quantization didn’t apply. Also check the execution_provider field — it should say CPUExecutionProvider or TensorrtExecutionProvider, not DnnlExecutionProvider (which doesn’t support INT8 well).

Pick Your Poison

If you’re shipping a pruned YOLOv8 model to edge devices, re-quantizing after pruning (Fix 1) is the safest bet. It’s fast, works with any pruning library, and ONNX Runtime handles it cleanly.

QAT-aware pruning (Fix 2) squeezes out an extra 2% mAP but costs you training time and requires deeper integration with your pruning pipeline. Worth it if you’re deploying to thousands of devices and every 0.01 mAP matters.

Manual graph surgery (Fix 3) is a last resort — use it to debug broken models or patch third-party ONNX files you can’t retrain. Don’t ship it to production unless you’ve validated the scale/zero-point values on real data.

One thing I haven’t tested yet: combining iterative magnitude pruning (prune 10% → fine-tune → prune another 10% → fine-tune) with progressive INT8 quantization. In theory, gradual quantization should let the model adapt to both sparsity and reduced precision simultaneously, but I suspect the training time would be prohibitive. If anyone’s tried this on YOLO-scale models, I’d love to see the numbers.

Debugging this kind of failure at 2am? Dark Chocolate Espresso Beans got me through more ONNX graph dumps than I care to admit.