Model-Quantization

A side-by-side comparison of eight neural-network quantization and precision strategies on a single convolutional image classifier. The pipeline trains one FP32 baseline once, then derives every other variant from it (or from scratch, where relevant) and reports accuracy, precision, F1, inference latency, model size, and training cost in one table.

src/main.py runs all eight stages in order and prints a final comparison table:

Stages 5-7 use PyTorch's fbgemm backend, which only ships CPU int8 kernels, so converted models are evaluated on CPU. BF16 is skipped automatically on GPUs below SM 8.0 (Ampere).

• Architecture — Fusable_Simple_CNN in src/model.py: three Conv2d → ReLU → MaxPool blocks followed by two linear layers, with QuantStub/DeQuantStub inserted so the same model can be quantized. The four Conv-ReLU / Linear-ReLU pairs are fused before quantization.
• Dataset — Caltech-101 (downloaded via torchvision), resized to 128×128, split 80/20 into train/validation, with a 200-image calibration subset carved out of the training set for static PTQ.

src/
├── config.py        # Single source of truth for all tunables (see below)
├── model.py         # Fusable_Simple_CNN architecture
├── data_loader.py   # Caltech-101 loaders (train / val / calibration)
├── train.py         # FP32 baseline training loop
├── ptq.py           # Post-training quantization (static + dynamic)
├── qat.py           # Quantization-aware training
├── quant_4bit.py    # Simulated 4-bit weight quantization + packing
├── evaluate.py      # Accuracy / precision / F1 and inference-time measurement
├── model_utils.py   # Save / load checkpoints, file-size helpers
├── plotting.py      # Comparison figures written to figures/
└── main.py          # Orchestrates all 8 stages and the final table

Outputs land in:

• models/ — trained checkpoints (fp32.pth, fp32_cpu.pth, fp16.pth, bf16.pth, ptq_dynamic.pth, ptq_static.pth, qat.pth, quant_4bit.pth)
• figures/ — comparison plots generated from the results

This project keeps every tunable value out of the training and quantization code and reads it from one place: src/config.py. It exposes a singleton config instance built from grouped dataclasses:

• TrainingHyperparameters — batch size, epochs, learning rate, scheduler, gradient clipping
• DataConfig — dataset path, validation split, calibration sample count, image size, normalization stats, DataLoader worker settings
• ModelConfig — layer sizes, kernel/padding, and the fuse_modules groups used by PTQ/QAT
• PTQConfig / QATConfig — quantization backend (fbgemm) and dynamic-quant layer types
• Float16Config / Quant4BitConfig — half-precision and 4-bit options
• EvalConfig — number of timed batches, warmup, averaging mode
• PathConfig / PlotConfig — output directories and figure format

Change a setting once in src/config.py and every stage picks it up. The Config class also provides get_device() (CUDA when available), get_quant_device() (always CPU, for int8 inference), and bf16_supported() (SM 8.0+ check).

pip install -r requirements.txt

The pinned requirements.txt targets a CUDA 12.4 build of PyTorch (torch==2.6.0+cu124, torchvision==0.21.0+cu124). On a CPU-only machine, replace those three lines with the matching CPU wheels from the PyTorch index.

Run the full comparison:

python src/main.py

The script prints per-stage progress, then a final table comparing accuracy, precision, F1, single-image and batched inference time, model size, and training time across all eight methods, and writes comparison figures into figures/.

To train only the FP32 baseline (useful for iteration):

python src/train.py

PyTorch has no native 4-bit inference kernels for CNNs, so src/quant_4bit.py implements a research-grade simulation: weights are quantized to 4-bit (15 symmetric levels, -7…+7, plus a per-tensor scale), packed two values per byte for storage, and dequantized back to FP32 for the actual forward pass. The file saved to disk is genuinely 4-bit packed (real size savings, roughly 4× smaller than FP32), and the accuracy measurement is valid because the model sees the same quantization noise it would with real 4-bit kernels — but there is no inference speedup, since the math runs as dequantized FP32.

The complete, runnable code for every stage lives in this repository under src/. Full attribution for every dataset and library used is in RESOURCES.md.