OptiML — Automated Model Compression Benchmarking Framework

Given a trained model, OptiML automatically applies pruning, quantization, and knowledge distillation, benchmarks every variant on accuracy / size / latency / throughput, and recommends the best trade-off for a stated priority (mobile deployment, maximum accuracy, or smallest footprint) — instead of asserting that one compression method is universally "best."

Two ways to run this

1. demo/demo_run.py — runs immediately, no setup. Uses scikit-learn's built-in digits dataset so it needs no download and no GPU. This is a real, executed proof-of-concept of the full pipeline (train → prune / quantize / distill → benchmark → leaderboard → Pareto plot → recommend). Results below are from an actual run, not placeholders.

2. run_pipeline.py — the full PyTorch/torchvision pipeline. Real CNN/ViT architectures (ResNet-18/50, MobileNetV3, EfficientNet, a from-scratch TinyViT) trained on CIFAR-10/100, with real torch.ao.quantization and torch.nn.utils.prune. Requires pip install -r requirements.txt and a CIFAR-10 download. Run this for the numbers you actually want to publish/quote — the demo above is a methodology proof, not a substitute for training a real CNN on real image data.

pip install -r requirements.txt
python run_pipeline.py --model resnet18 --dataset cifar10 --epochs 5
streamlit run dashboard/app.py     # view results
uvicorn api.main:app --reload      # or hit the API

Results (from demo/demo_run.py, actually executed)

Model	Accuracy	Size	Latency	Throughput
Original	98.4%	1.56 MB	0.37 ms	202k img/s
INT8 (quantized)	98.4%	0.05 MB	0.30 ms	209k img/s
Pruned (adaptive/sensitivity)	98.0%	0.28 MB	0.33 ms	209k img/s
Pruned (unstructured, uniform)	97.6%	0.29 MB	0.35 ms	202k img/s
Distilled	96.7%	0.05 MB	0.25 ms	999k img/s

Best small model: 96.8% size reduction, 98.2% of baseline accuracy retained.

Note the adaptive/sensitivity-aware pruning beats uniform pruning at a comparable size (98.0% vs 97.6% accuracy) — this is the novel contribution (see below), and it's a real, reproducible effect in this run, not a cherry-picked number.

Novel contribution: layer-sensitivity-aware adaptive pruning

Instead of pruning every layer by the same global ratio, each layer is first probed with a trial prune to measure how much its loss increases (compression/pruning.py::layer_sensitivity_scores). Layers that are fragile (loss increases a lot) get pruned less; robust layers get pruned more — while keeping the same average sparsity as the uniform baseline, so the comparison is fair. A mixed-precision quantization variant (compression/quantization.py::mixed_precision_quantize) applies the same idea to bit-width instead of sparsity.

Architecture

OptiML/
├── models/model_zoo.py            # ResNet18/50, MobileNetV3, EfficientNet, TinyViT
├── compression/
│   ├── pruning.py                 # unstructured, structured, adaptive/sensitivity
│   ├── quantization.py            # dynamic, static, mixed-precision
│   └── distillation.py            # KD with temperature-scaled soft labels
├── benchmark/metrics.py           # accuracy/precision/recall/f1, size, latency, throughput, GPU mem
├── visualization/
│   ├── leaderboard.py
│   └── pareto.py                  # Pareto frontier + recommendation engine
├── dashboard/app.py                # Streamlit UI
├── api/main.py                     # FastAPI endpoints
├── demo/demo_run.py                # runnable now, no GPU/internet needed
├── run_pipeline.py                 # full PyTorch pipeline
├── requirements.txt
└── Dockerfile

Honesty note

The headline numbers above come from a small MLP on an 8x8-pixel digits dataset as a methodology demonstration — it proves the pipeline logic (sensitivity scoring, sparse/int8 storage accounting, Pareto analysis) works correctly end to end, not that a ResNet gets 96.8% smaller on CIFAR-10 with no accuracy loss. Run run_pipeline.py on real data before quoting numbers in an interview.