Quantize Callback

Quantize your network during training

Overview

The QuantizeCallback enables Quantization-Aware Training (QAT) within the fastai training loop. QAT simulates quantization effects during training, allowing the model to adapt its weights for better accuracy after quantization.

Why use QAT over post-training quantization?

QAT aims to recover the accuracy the post-training path loses, by letting the weights adapt to the rounding they will be subject to. The effect is model-dependent — measure it on yours.

Trade-offs: - Requires retraining (not just calibration) - Every step carries the simulated quantization, so training costs more than the float run - Only for situations where you can afford additional training time

Two flows are available, and the backend picks between them: the FX backends ('x86', 'qnnpack', 'fbgemm', 'onednn') prepare the model with prepare_qat_fx, while 'pt2e' captures it with torch.export and prepares it with prepare_qat_pt2e — the same symmetric INT8 precision the post-training pt2e path applies, so the trained model can be written as a QDQ ONNX graph.

Preparation hands back a different module, so the callback rebuilds learn.opt on it, carrying over the hypers fit was called with and the per-parameter marks fastai sets by walking the model. Splitters written for the source model — every fastai vision splitter indexes it, and a prepared graph is not indexable — cannot be applied to the prepared one, so the groups are then rebuilt by parameter identity: the discriminative learning rates keep working, and parameters the preparation added, such as learnable fake-quantize scales, join the last group, and so train at the head’s learning rate.

source

QuantizeCallback

def QuantizeCallback(
    quantizer:fasterai.quantize.quantizer.Quantizer | None=None, # Custom quantizer; its backend decides which flow runs
    backend:str='x86', # Target backend: 'x86', 'qnnpack', 'fbgemm', 'onednn' or 'pt2e'
    use_per_tensor:bool=False, # Force per-tensor quantization (FX backends only)
    verbose:bool=False
):

Quantization-Aware Training (QAT), with the FX backends or with the pt2e flow. before_fit swaps in a prepared model carrying fake-quantize modules, after_fit converts it.

Parameters:

  • backend: only used if quantizer is not provided

Usage Example

from fasterai.quantize.quantize_callback import QuantizeCallback

# Basic QAT with default settings
cb = QuantizeCallback(backend='x86', verbose=True)

# Train with QAT
learn.fit(5, cbs=[cb])

# After training, the quantized model is available at:
quantized_model = learn.quantized_model

QAT Workflow

  1. before_fit: Model is prepared for QAT (fake quantization nodes inserted), and the optimizer is rebound to it — Learner.fit builds the optimizer before callbacks run, on the model as it was
  2. Training: Model trains with simulated quantization effects
  3. after_fit: Model is converted to fully quantized form

The final learn.model is the quantized model ready for CPU inference.

Choosing a backend

Backend What it produces Use it for
'x86', 'fbgemm', 'onednn', 'qnnpack' (FX) An FX quantized model, activations affine CPU inference through PyTorch
'pt2e' A graph captured by torch.export, symmetric INT8 everywhere A QDQ ONNX export (export_qdq) 
from fasterai.export.onnx_exporter import export_qdq

# QAT through the pt2e flow, then export the trained INT8 graph
learn.fit(5, cbs=[QuantizeCallback(backend='pt2e')])
export_qdq(learn.model, sample, 'model_qat.onnx')

The pt2e flow captures the model with torch.export, which specialises the graph to one batch size: every batch of both dataloaders must have that size, and the callback refuses dataloaders that do not. Any precision the grammar names can be run by passing the quantizer instead of the backend, e.g. QuantizeCallback(quantizer=Quantizer(backend='pt2e', method='qat', qscheme='per_tensor')).

A quantizer handed in this way keeps the method it was built with, so spell method='qat' — a quantizer built without it would run the post-training preparation inside the fit:

from fasterai.quantize.quantizer import Quantizer

# QAT that trains THROUGH the Q/DQ placement it will be exported with
cb = QuantizeCallback(quantizer=Quantizer(backend='pt2e', method='qat',
                                          qdq_placement='skip_conv_add'))

qdq_placement='skip_conv_add' leaves the residual branch’s convolution reaching its addition without a Q/DQ pair. Under QAT the model trains through that placement, so the weights adapt to the arithmetic that will actually run — and, unlike the post-training path, the scales the graph keeps are not the ones the ordinary placement would have observed: removing a fake-quantize module changes every statistic downstream of it.

See Also

Tests live in nbs/tests/test_quantize_callback.ipynb.