Quantization aware training for a small CNN

This repo shows how to take a small convolutional network from float32 weights down to int8 using quantization aware training (QAT) in PyTorch. Instead of quantizing a trained model after the fact, QAT inserts fake quantization into the forward pass so the network feels the rounding error during training and adjusts its weights to compensate. When you then convert to a real int8 model, the accuracy you keep is usually better than plain post training quantization.

What is inside

src/model.py holds SmallCNN, a compact two block conv net for tiny grayscale images such as the MNIST family. It wraps its forward pass in a QuantStub and DeQuantStub so the float and quantized regions are explicit, and it knows how to fuse its conv, batch norm and ReLU layers, which is what lets the int8 conversion behave well.

src/qat.py is the pipeline. The important pieces are:

• build_qat_model fuses the float model, attaches a QAT qconfig and runs prepare_qat, returning a model with fake quantization observers in place.
• qat_train_step and train_qat run ordinary training. The fake quant modules make the model simulate int8 rounding while still computing in float, so gradients flow normally.
• convert_to_int8 turns the trained QAT model into a real int8 model whose weights are stored as qint8 and whose layers run the quantized kernels.
• make_synthetic_loader builds a small random dataset with a faint per class signal, which keeps the tests fast and free of any download.

The quantization engine is chosen at runtime. Server wheels usually expose fbgemm, but this machine only ships onednn, so default_backend picks whatever the build reports rather than hardcoding one name.

Running it

Install the dependencies and run the demo:

pip install -r requirements.txt
python demo.py

On the machine where this was written the demo trained for three short epochs on synthetic data, converted to int8 on the onednn engine, and the int8 model tracked the fake quantized float model closely. The largest absolute difference in the raw logits across a batch of sixteen inputs was about 0.023, and the two models agreed on the argmax class for roughly 94 percent of those inputs. Your numbers will move a little with the random seed and the engine, but the int8 model should stay close to the float reference.

Tests

pytest tests/ -q

The suite checks behavior rather than fixed numbers:

• the float model produces finite logits of the right shape,
• preparing for QAT actually inserts fake quantization modules,
• a single training step runs and moves the weights,
• conversion yields genuine int8 quantized weights,
• the int8 output stays within tolerance of the fake quantized float output on synthetic input, and
• the two models agree on the predicted class for most inputs.

All six tests pass on CPU with no downloads.

A note on the reference model

The tests compare the int8 model against the QAT model in eval mode with fake quantization still active, not against the original float32 model. That is deliberate. Fake quantization is the function the int8 kernels are built to reproduce, so it is the honest target. Comparing straight to the unquantized float model would fold in the quantization error itself, which is expected and not a bug.