An exploratory project implementing neural network quantization (and, eventually, pruning) from scratch, as a way to build a solid, first-principles understanding of model compression techniques used in embedded/edge ML deployment.
The guiding idea: rather than calling torch.quantization.quantize_dynamic() as a black box, this project derives the underlying math by hand and implements the quantization pipeline manually in C++, then verifies the impact on a trained model's accuracy in Python.
Res8 is a small residual CNN for keyword spotting, trained here on Google Speech Commands v0.02.
- Baseline accuracy: ~91% (float32, full precision)
- Task: classify short audio clips into command words / "unknown" / "background noise"
This model was chosen because it's small enough to fully inspect layer-by-layer, yet realistic enough (real conv layers, BatchNorm, residual blocks) to analyze quantization challenges.
cpp/ C++ tensor + quantization library ("cmc")
├── include/cmc/ Public headers (tensor.hpp, io.hpp, quantization.hpp, metrics.hpp)
├── src/ Implementations
└── CMakeLists.txt
data/
├── raw/ Speech Commands dataset
├── checkpoints/ Trained PyTorch checkpoints (.pth)
└── exported/ Binary weight files exchanged between Python and C++
python/
├── model.py Res8 architecture
├── dataset.py Speech Commands dataset loader
├── train.py / tune.py Training pipeline
└── compression/
├── model_io.py Export/import weights to/from the custom binary format
└── evaluate.py Runs the validation loop on a (re)imported model
To move weights between PyTorch and C++, checkpoints are serialized into a small, self-describing binary format, not tied to any external serialization library:
[magic: 4 bytes, "RES8"]
[version: uint32]
[num_tensors: uint32]
for each tensor:
[name_len: uint32][name: name_len bytes, UTF-8]
[ndim: uint32][shape: ndim x uint32]
[data: prod(shape) x float32]
This format is written and read symmetrically on both sides:
- Python → binary:
compression/model_io.py::export_state_dict - binary → C++:
cmc::load_weights - C++ → binary:
cmc::write_weights - binary → Python:
compression/model_io.py::import_state_dict
The C++ library (cmc) currently implements symmetric, per-tensor, post-training quantization (PTQ) to INT8:
- Scale is computed from the tensor's maximum absolute value:
s = max(|x|) / 127 - Zero-point is fixed at
z = 0(a consequence of the symmetry assumption, not an independent choice) - The range
[-128, 127]ofint8is deliberately restricted to[-127, 127]to keep the scale consistent on both sides of zero
The full derivation of these formulas, starting from the general affine quantization mapping and specializing step by step to the symmetric case, is available in docs/calculations.md.
For verification purposes, weights are represented as quantize → dequantize round-trips on float32 tensors (a form of "fake quantization"), rather than as physical int8 buffers. This is a deliberate simplification. The goal at this stage is to validate the math of quantization and its effect on model accuracy, not to build a full integer-only inference runtime (which would additionally require activation quantization, requantization between layers, and BatchNorm folding — see Future Work).
1. train.py → data/checkpoints/best_model.pth
2. model_io.export_state_dict → data/exported/weights.bin
3. cmc::load_weights (C++) → in-memory Tensor objects
4. cmc::quantize / dequantize → reconstructed Tensor objects
5. cmc::write_weights (C++) → data/exported/reconstructed_weights.bin
6. model_io.import_state_dict → PyTorch state_dict
7. compression/evaluate.py → validation accuracy
# 1. Train (or use an existing checkpoint)
cd python && python train.py
# 2. Export weights to the binary format
python -m compression.model_io # or however export is currently triggered
# 3. Build and run the C++ quantization pipeline
cd ../cpp && mkdir -p build && cd build
cmake .. && cmake --build .
./cmc_main
# 4. Evaluate the reconstructed model
cd ../../python
python -m compression.evaluate| Metric | Value |
|---|---|
| Baseline accuracy (float32) | ~91.0% |
| Accuracy after symmetric INT8 quantize→dequantize (all layers) | 89.56% |
| Accuracy drop | ~1.5 percentage points |
Measuring per-tensor MSE between original and reconstructed weights revealed a consistent pattern. BatchNorm's running_var reconstructs noticeably worse than convolutional weights, with roughly 3-4 orders of magnitude higher MSE within the same block. Convolutional weights are typically small and near-symmetric around zero, a natural consequence of initialization and regularization. running_var has a wider, non-negative range that is more sensitive to outliers. Since symmetric quantization derives its scale from the single largest-magnitude value in the tensor, an unusually large value in running_var disproportionately widens the quantization step for the entire tensor.
This is consistent with the general finding (see Nagel et al., 2021) that simple min-max range setting is more outlier-sensitive than alternatives like MSE-based range setting, particularly for tensors with skewed or heavy-tailed distributions.
num_batches_tracked buffers were found to reconstruct with zero error (as expected for small integers on a coarse quantization grid) and are not meaningful to quantize. They should ideally be excluded from the pipeline entirely.
This project intentionally stops short of production-grade quantization. Known gaps:
- Only symmetric, per-tensor quantization is implemented. Per-channel quantization (especially for weights) and MSE-based or cross-entropy-based range setting are known to outperform simple min-max/symmetric schemes, particularly at lower bit-widths.
- Activations are not quantized. Only weights are currently affected. A full PTQ pipeline requires calibrating activation ranges on representative data as well.
- No BatchNorm folding. In a real deployment pipeline, BatchNorm parameters are algebraically folded into the preceding convolution before quantization, rather than quantized as standalone tensors. This would likely eliminate the
running_varerror pattern described above. - No integer-only inference. The current pipeline verifies accuracy via simulated "fake" quantization on float32 tensors. A true integer-only runtime would require int8×int8→int32 accumulation, requantization between layers via fixed-point multipliers, and integer-friendly handling of nonlinearities, a substantially larger undertaking.
- Pruning is not yet implemented. Planned as the next major addition to this project, alongside a joint pruning+quantization comparison.
- Nagel, M. et al. (2021). A White Paper on Neural Network Quantization. arXiv:2106.08295
- Jacob, B. et al. (2017). Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. arXiv:1712.05877
- Han, S. et al. (2015). Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. arXiv:1510.00149
- Tang, R. & Lin, J. (2018). Deep Residual Learning for Small-Footprint Keyword Spotting. arXiv:1710.10361
- Warden, P. (2018). Speech Commands: A Dataset for Limited-Vocabulary Speech Recognition. arXiv:1804.03209