model-compression

Two classic post training compression techniques for a small PyTorch model: magnitude based weight pruning and dynamic int8 quantization. The code is deliberately compact so the whole thing runs on a CPU in about a second with no downloads. Everything is covered by property style unit tests.

What is here

The model is a three layer perceptron defined in src/model.py. It is small but real, with enough parameters that pruning and quantization actually change the weight tensors in a measurable way.

src/pruning.py implements unstructured magnitude pruning. Weights with the smallest absolute value get zeroed until a requested fraction of the tensor is sparse. Two scopes are available. Per layer pruning drives every weight tensor to the same target sparsity on its own. Global pruning computes a single threshold across all prunable weights at once, so layers with larger weights keep more of them and the overall sparsity matches the request. One detail worth calling out: when several weights sit exactly at the cutoff magnitude, the code trims the surplus deterministically so the achieved sparsity is exact rather than off by a few weights.

src/quantization.py wraps PyTorch dynamic quantization. Linear layers are stored as int8 and activations are quantized at inference time, so no calibration data is needed. The original model is deep copied first, so quantizing never mutates the float model you passed in.

Layout

src/
  model.py          small MLP and a deterministic builder
  pruning.py        magnitude pruning, per layer and global, plus sparsity meters
  quantization.py   dynamic int8 quantization and output comparison helpers
tests/
  test_pruning.py
  test_quantization.py

Usage

import torch
from src.model import build_model
from src.pruning import prune_model, sparsity
from src.quantization import quantize_dynamic, max_abs_output_diff

model = build_model(seed=0)

# Prune half of the weights using a single global threshold.
prune_model(model, amount=0.5, global_prune=True)
print("sparsity:", sparsity(model))

# Then apply dynamic int8 quantization on top.
qmodel = quantize_dynamic(model)

x = torch.randn(16, 64)
print("max output difference:", max_abs_output_diff(model, qmodel, x))

What the tests check

The tests are behavior checks rather than smoke tests. For pruning they confirm that the requested sparsity is hit exactly across a range of amounts, that the weights dropped really are the smallest by magnitude, that a global threshold prunes a tiny weight layer harder than a large weight layer while still matching the overall target, and that pruning with amount zero leaves outputs unchanged. For quantization they confirm that every Linear layer converts, that the float model is left untouched, that the quantized output stays close to the original on random input within a tolerance, and that the full prune then quantize pipeline still tracks the pruned baseline.

Running

pip install -r requirements.txt
pytest tests/ -q

On the machine where this was built, the suite reported 19 passed.

Notes

Dynamic quantization here targets Linear layers, which is where this model keeps all of its parameters. The pruning masks are written straight into the weights as zeros, so a later quantization step sees the sparse tensor and the two techniques compose cleanly.