04: Model Compression (Pruning & Quantization)¶

This notebook will illustrate every model compression technique provided by shrinkai:

• Pruner, masking-based pruning (unstructured and structured)
• ChannelPruner, physically removing channels from Conv2d/Linear layers
• Quantizer, Post-Training Quantization (dynamic and static) and Quantization-Aware Training
• How to combine pruning, quantization, and distillation together

Unlike distillation, compression does not need a second (teacher) model. It works directly on a single trained network. The use case is the same CIFAR-10 setup as the first tutorial, using small pretrained CIFAR-10 models. As always, the goal is to show how to use shrinkai correctly and understand why each technique behaves the way it does, not to reach state-of-the-art performance.

Setup¶

import copy

import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

from shrinkai.profiler import count_parameters, estimate_model_size_mb, measure_latency, count_flops

import copy import torch from torch.utils.data import DataLoader, Subset from torchvision import datasets, transforms from shrinkai.profiler import count_parameters, estimate_model_size_mb, measure_latency, count_flops

We reuse the same lightweight resnet20 architecture that was the student in the first tutorial, but here it is loaded pretrained, as compression is applied to an already-trained model, it does not require training from scratch.

resnet = torch.hub.load(
    "chenyaofo/pytorch-cifar-models",
    "cifar10_resnet20",
    pretrained=True,
    verbose=False,
).eval()

resnet = torch.hub.load( "chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=True, verbose=False, ).eval()

def build_cifar_dataloaders(
    batch_size: int = 64,
    subset_size: int | None = 500,
) -> tuple[DataLoader, DataLoader, torch.Tensor]:
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
        ]
    )

    train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
    val_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)

    if subset_size is not None:
        train_dataset = Subset(train_dataset, range(subset_size))
        val_dataset = Subset(val_dataset, range(subset_size))

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

    sample_input = torch.randn(batch_size, 3, 32, 32)

    return train_loader, val_loader, sample_input

def build_cifar_dataloaders( batch_size: int = 64, subset_size: int | None = 500, ) -> tuple[DataLoader, DataLoader, torch.Tensor]: transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ] ) train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) val_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) if subset_size is not None: train_dataset = Subset(train_dataset, range(subset_size)) val_dataset = Subset(val_dataset, range(subset_size)) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2) sample_input = torch.randn(batch_size, 3, 32, 32) return train_loader, val_loader, sample_input

train_loader, val_loader, sample_input = build_cifar_dataloaders(batch_size=64, subset_size=500)

train_loader, val_loader, sample_input = build_cifar_dataloaders(batch_size=64, subset_size=500)

Before touching the model, let's establish a baseline using shrinkai.profiler directly. The same building blocks that Profiler.compare() (used by every .benchmark() method below) aggregates internally:

• count_parameters: the raw number of learnable weights (including biases).
• estimate_model_size_mb: the actual size of model.state_dict() once serialized, i.e. what you'd ship to a device.
• measure_latency: wall-clock inference time, averaged over many runs after a warmup (hardware-dependent).
• count_flops: the number of floating-point operations for one forward pass, hardware-independent, via torch.utils.flop_counter.

These four numbers do not always move together, which is precisely the point of walking through every technique below: a method can shrink one of them a lot while leaving the others almost untouched.

print("Parameters:", count_parameters(resnet)["total_params"])
print("Disk size (MB):", estimate_model_size_mb(resnet))
print("FLOPs / sample:", count_flops(resnet, sample_input, device="cpu") // sample_input.size(0))
print(measure_latency(resnet, sample_input, torch.device("cpu"), num_runs=30, warmup_runs=5))

print("Parameters:", count_parameters(resnet)["total_params"]) print("Disk size (MB):", estimate_model_size_mb(resnet)) print("FLOPs / sample:", count_flops(resnet, sample_input, device="cpu") // sample_input.size(0)) print(measure_latency(resnet, sample_input, torch.device("cpu"), num_runs=30, warmup_runs=5))

Parameters: 272474
Disk size (MB): 1.08
FLOPs / sample: 81626368
{'batch_latency_ms': 65.51553467094588, 'sample_latency_ms': 1.0236802292335294, 'fps': 976.8675524277149}

Pruning¶

Pruning removes parameters that contribute little to a model's output, based on the idea that trained networks are heavily over-parameterized. In this package, there are 2 different mechanisms:

• Pruner: attaches a mask to each targeted layer's weights (via torch.nn.utils.prune). The tensor shape never changes, but pruned weights are zeroed out but still physically stored and still multiplied at inference time.
• ChannelPruner: physically removes channels from Conv2d/Linear layers (and the BatchNorm that depends on them), actually shrinking the tensors.

This distinction matters enormously in practice, and the first part of this section exists specifically to make that gap tangible with real numbers, before introducing ChannelPruner.

Unstructured pruning¶

L1-unstructured pruning ranks individual weights by absolute magnitude and zeroes out the smallest fraction:

$$\text{mask}_i = \begin{cases} 0 & \text{if } |w_i| \text{ is among the } k \text{ smallest magnitudes} \\ 1 & \text{otherwise} \end{cases}, \qquad k = \lfloor \text{amount} \times n \rfloor$$

The intuition (going back to LeCun et al.'s Optimal Brain Damage, 1989) is that a weight close to zero contributes little to the overall function the network computes, so removing it should barely affect predictions, but the pattern of zeros is essentially random from a hardware point of view (any weight, anywhere in the tensor, can be zeroed), which will matter in a moment.

from shrinkai.compression.pruning import Pruner, PruningConfig

from shrinkai.compression.pruning import Pruner, PruningConfig

config = PruningConfig(method="unstructured", amount=0.3)
pruner = Pruner(config)

unstructured_pruned = pruner.apply(copy.deepcopy(resnet))
print("fc has a weight_mask buffer:", hasattr(unstructured_pruned.fc, "weight_mask"))
print("zeroed weights in fc:", (unstructured_pruned.fc.weight == 0).sum().item(), "/", unstructured_pruned.fc.weight.numel())

config = PruningConfig(method="unstructured", amount=0.3) pruner = Pruner(config) unstructured_pruned = pruner.apply(copy.deepcopy(resnet)) print("fc has a weight_mask buffer:", hasattr(unstructured_pruned.fc, "weight_mask")) print("zeroed weights in fc:", (unstructured_pruned.fc.weight == 0).sum().item(), "/", unstructured_pruned.fc.weight.numel())

fc has a weight_mask buffer: True
zeroed weights in fc: 192 / 640

Pruner.apply() attaches a forward pre-hook per layer: at each forward pass, PyTorch recomputes weight = weight_orig * weight_mask. This is what makes pruning compatible with pruning-aware training: you can keep training (or distilling) the model afterward, gradients only flow to the unpruned weights, and the mask stays enforced at every step. Before deploying the model, call Pruner.finalize() once to bake the mask permanently into the weights and remove the hook.

report = pruner.benchmark(
    resnet,
    unstructured_pruned,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

report = pruner.benchmark( resnet, unstructured_pruned, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

                                          Benchmark Report                                           
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (Dense)) ┃ Student (Pruned (Sparse)) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                     0.27 M │                    0.27 M │                -0.0% │
│ Model Size (Disk) │                    1.08 MB │                   1.08 MB │ -0.0% (1.0x smaller) │
│ Latency / Sample  │                    0.95 ms │                   0.96 ms │          1.0x faster │
│ Throughput (FPS)  │               1049.6 img/s │              1043.0 img/s │   +1.0x (1043.0 FPS) │
│ Accuracy          │                     91.80% │                    90.40% │       98.5% retained │
└───────────────────┴────────────────────────────┴───────────────────────────┴──────────────────────┘

The number that matters here is "Model Size" and "Latency", and they barely moved. 30% of fc's weights (and every other targeted layer) are now exactly zero, yet the file is the same size on disk and inference is not faster. This is not a bug, it is the defining limitation of mask-based pruning: a saved tensor with torch.save stores every value, zero or not, at full precision, and a dense matrix multiply still visits every zeroed cell. Realizing an actual speedup from unstructured sparsity requires a sparse tensor format and a sparse-aware compute kernel, neither of which vanilla Pruner sets up (some accelerators / server GPUs with structured 2:4 sparsity support, dedicated sparse inference runtimes can exploit it, but that is a deployment-side capability, not something Pruner itself provides). The one number that did move a bit is accuracy: pruning is a lossy transformation, and 30% of the weakest weights still cost a little of it.

Structured pruning¶

Structured pruning ranks entire channels/filters instead of individual weights, using their $L_n$-norm (here $n=2$):

$$\|\mathbf{w}_{c}\|_2 = \sqrt{\sum_j w_{c,j}^2}, \qquad \text{drop the channels } c \text{ with the smallest norm}$$

The appeal is that removing a whole channel is, in principle, a step toward a smaller dense computation, but Pruner's "structured" method still only masks the channel (sets every weight in it to zero); it does not resize the tensor. So it inherits the exact same "no real speedup" limitation as unstructured pruning, with one additional risk: because it commits to removing whole channels rather than the weakest individual scalars, it is a much coarser, all-or-nothing decision. Let's see what that costs on a small model with virtually no redundancy to spare.

config = PruningConfig(method="structured", amount=0.2)
pruner = Pruner(config)

structured_pruned = pruner.apply(copy.deepcopy(resnet))
report = pruner.benchmark(
    resnet,
    structured_pruned,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

config = PruningConfig(method="structured", amount=0.2) pruner = Pruner(config) structured_pruned = pruner.apply(copy.deepcopy(resnet)) report = pruner.benchmark( resnet, structured_pruned, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

                                          Benchmark Report                                           
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (Dense)) ┃ Student (Pruned (Sparse)) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                     0.27 M │                    0.27 M │                -0.0% │
│ Model Size (Disk) │                    1.08 MB │                   1.08 MB │ -0.0% (1.0x smaller) │
│ Latency / Sample  │                    0.96 ms │                   0.96 ms │          1.0x faster │
│ Throughput (FPS)  │               1046.8 img/s │              1041.1 img/s │   +1.0x (1041.1 FPS) │
│ Accuracy          │                     91.80% │                     8.00% │        8.7% retained │
└───────────────────┴────────────────────────────┴───────────────────────────┴──────────────────────┘

Accuracy collapses almost to random guessing (10% on 10 classes), for only 20% pruning. Two things compound here: resnet20 is a genuinely tiny network (272K parameters, as few as 16 channels in its first stage) with very little redundancy to begin with, and Pruner(method="structured") prunes every targeted layer (Linear and Conv2d by default) simultaneously, all at once, with no fine-tuning step in between, the damage compounds across all ~20 convolutions in the network at the same time. Larger, more over-parameterized models tolerate structured pruning far better, and in practice you would always fine-tune afterward (e.g. resume training through a Distiller, using the pruned model as the student and the original as the teacher) to recover the accuracy lost by the pruning step. We still have not gained any real speedup, though, for that, we need to talk about physically removing the channels rather than masking them.

Channel pruning¶

ChannelPruner rebuilds smaller Conv2d/Linear layers with the pruned channels actually removed, instead of masked. Doing this correctly is harder than it looks: removing output channel $c$ from a layer is only valid if every downstream consumer of that output (next Conv2d/Linear's input dimension or any BatchNorm sitting in between) shrinks to match. ChannelPruner uses torch.fx to trace the model's real dataflow graph (not just its list of named modules) and walks it forward from each target layer to find exactly what needs to shrink alongside it.

By design, it only proceeds when that walk is unambiguous: the target layer's output must feed, through any number of BatchNorm/activation/pooling layers, into exactly one consumer (another Conv2d/Linear, or the model's own output). Anything else, such as a skip connection, a concatenation, a grouped/depthwise convolution, is rejected with a clear error rather than silently producing an incorrect model. Let's look at resnet20's structure to find layers where this is safe.

print(dict(resnet.named_modules()).keys())

print(dict(resnet.named_modules()).keys())

dict_keys(['', 'conv1', 'bn1', 'relu', 'layer1', 'layer1.0', 'layer1.0.conv1', 'layer1.0.bn1', 'layer1.0.relu', 'layer1.0.conv2', 'layer1.0.bn2', 'layer1.1', 'layer1.1.conv1', 'layer1.1.bn1', 'layer1.1.relu', 'layer1.1.conv2', 'layer1.1.bn2', 'layer1.2', 'layer1.2.conv1', 'layer1.2.bn1', 'layer1.2.relu', 'layer1.2.conv2', 'layer1.2.bn2', 'layer2', 'layer2.0', 'layer2.0.conv1', 'layer2.0.bn1', 'layer2.0.relu', 'layer2.0.conv2', 'layer2.0.bn2', 'layer2.0.downsample', 'layer2.0.downsample.0', 'layer2.0.downsample.1', 'layer2.1', 'layer2.1.conv1', 'layer2.1.bn1', 'layer2.1.relu', 'layer2.1.conv2', 'layer2.1.bn2', 'layer2.2', 'layer2.2.conv1', 'layer2.2.bn1', 'layer2.2.relu', 'layer2.2.conv2', 'layer2.2.bn2', 'layer3', 'layer3.0', 'layer3.0.conv1', 'layer3.0.bn1', 'layer3.0.relu', 'layer3.0.conv2', 'layer3.0.bn2', 'layer3.0.downsample', 'layer3.0.downsample.0', 'layer3.0.downsample.1', 'layer3.1', 'layer3.1.conv1', 'layer3.1.bn1', 'layer3.1.relu', 'layer3.1.conv2', 'layer3.1.bn2', 'layer3.2', 'layer3.2.conv1', 'layer3.2.bn1', 'layer3.2.relu', 'layer3.2.conv2', 'layer3.2.bn2', 'avgpool', 'fc'])

Every BasicBlock in this network (e.g. layer1.0) is conv1 → bn1 → relu → conv2 → bn2, followed by adding the block's own input back in (the residual/skip connection) before moving to the next block. That means:

• layer1.0.conv1's output feeds only bn1 → relu → layer1.0.conv2, a single, unambiguous chain. Safe to prune.
• layer1.0.conv2's output feeds bn2, then the residual addition with the block's input, two tensors merging. Not safe, and ChannelPruner will refuse it.

Let's prune the first case.

from shrinkai.compression.pruning import ChannelPruner

from shrinkai.compression.pruning import ChannelPruner

channel_pruner = ChannelPruner(amount=0.2)
gently_pruned = channel_pruner.apply(
    copy.deepcopy(resnet),
    target_layers=["layer1.0.conv1"],
    sample_input=sample_input,
)

print("conv1 out_channels:", gently_pruned.layer1[0].conv1.out_channels, "(was 16)")
print("bn1 num_features:", gently_pruned.layer1[0].bn1.num_features)
print("conv2 in_channels:", gently_pruned.layer1[0].conv2.in_channels)

channel_pruner = ChannelPruner(amount=0.2) gently_pruned = channel_pruner.apply( copy.deepcopy(resnet), target_layers=["layer1.0.conv1"], sample_input=sample_input, ) print("conv1 out_channels:", gently_pruned.layer1[0].conv1.out_channels, "(was 16)") print("bn1 num_features:", gently_pruned.layer1[0].bn1.num_features) print("conv2 in_channels:", gently_pruned.layer1[0].conv2.in_channels)

conv1 out_channels: 13 (was 16)
bn1 num_features: 13
conv2 in_channels: 13

report = channel_pruner.benchmark(
    resnet,
    gently_pruned,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
    compute_flops=True,
)
report.show()

report = channel_pruner.benchmark( resnet, gently_pruned, sample_input=sample_input, val_dataloader=val_loader, device="cpu", compute_flops=True, ) report.show()

                                                Benchmark Report                                                
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (Dense)) ┃ Student (Pruned (Physically Shrunk)) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                     0.27 M │                               0.27 M │                -0.3% │
│ Model Size (Disk) │                    1.08 MB │                              1.08 MB │ -0.0% (1.0x smaller) │
│ Latency / Sample  │                    0.98 ms │                              0.98 ms │          1.0x faster │
│ Throughput (FPS)  │               1015.8 img/s │                         1015.4 img/s │   +1.0x (1015.4 FPS) │
│ FLOPs / Sample    │             5224.09 MFLOPs │                       5110.84 MFLOPs │                -2.2% │
│ Accuracy          │                     91.80% │                               91.80% │      100.0% retained │
└───────────────────┴────────────────────────────┴──────────────────────────────────────┴──────────────────────┘

This time, parameters, disk size, latency, and FLOPs all decreased, for zero accuracy cost, because we only removed the 20% least-useful channels of a single, carefully-chosen layer. Let's push it further across two layers to see the trade-off curve appear:

aggressively_pruned = ChannelPruner(amount=0.3).apply(
    copy.deepcopy(resnet),
    target_layers=["layer1.0.conv1", "layer2.0.conv1"],
    sample_input=sample_input,
)

report = ChannelPruner(amount=0.3).benchmark(
    resnet,
    aggressively_pruned,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
    compute_flops=True,
)
report.show()

aggressively_pruned = ChannelPruner(amount=0.3).apply( copy.deepcopy(resnet), target_layers=["layer1.0.conv1", "layer2.0.conv1"], sample_input=sample_input, ) report = ChannelPruner(amount=0.3).benchmark( resnet, aggressively_pruned, sample_input=sample_input, val_dataloader=val_loader, device="cpu", compute_flops=True, ) report.show()

                                                Benchmark Report                                                
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (Dense)) ┃ Student (Pruned (Physically Shrunk)) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                     0.27 M │                               0.27 M │                -1.9% │
│ Model Size (Disk) │                    1.08 MB │                              1.06 MB │ -1.9% (1.0x smaller) │
│ Latency / Sample  │                    0.96 ms │                              0.94 ms │          1.0x faster │
│ Throughput (FPS)  │               1041.1 img/s │                         1060.0 img/s │   +1.0x (1060.0 FPS) │
│ FLOPs / Sample    │             5224.09 MFLOPs │                       4945.69 MFLOPs │                -5.3% │
│ Accuracy          │                     91.80% │                               83.40% │       90.8% retained │
└───────────────────┴────────────────────────────┴──────────────────────────────────────┴──────────────────────┘

More channels pruned across more layers means a bigger real speedup, but this time at a visible accuracy cost, the same trade-off that appeared with masked structured pruning, except now it comes attached to gains you can actually deploy. In practice, this is exactly where you may distill: pass aggressively_pruned as the student to a Distiller, with resnet as the teacher, and recover most of the lost accuracy in a few epochs (see Tutorial 01). Physical pruning and distillation are complementary, not alternatives.

Let's confirm ChannelPruner actually refuses layer1.0.conv2 (the one feeding the residual addition), instead of silently producing a broken model:

try:
    ChannelPruner(amount=0.2).apply(copy.deepcopy(resnet), target_layers=["layer1.0.conv2"], sample_input=sample_input)
except ValueError as e:
    print(f"Rejected, as expected:\n{e}")

try: ChannelPruner(amount=0.2).apply(copy.deepcopy(resnet), target_layers=["layer1.0.conv2"], sample_input=sample_input) except ValueError as e: print(f"Rejected, as expected:\n{e}")

Rejected, as expected:
Layer 'layer1.0.conv2' output passes through unsupported call_function '<built-in function add>' before reaching a Conv2d/Linear layer. ChannelPruner cannot guarantee this is safe to prune through; use `Pruner` instead for this topology.

Pruning through a classifier head.

resnet20's residual blocks mean no convolution here ever feeds avgpool/fc directly without crossing a residual addition first, so that specific case cannot be demonstrated on this model. It is common enough to be worth seeing in isolation: as long as a Conv2d's output is pooled down to a single spatial position (e.g. via AdaptiveAvgPool2d(1)) before being flattened into a Linear layer, ChannelPruner can prune straight through into that Linear's input dimension too — because at that point each channel maps to exactly one flattened feature, with no risk of the interleaving ambiguity a larger spatial size would create.

import torch.nn as nn


class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)


tiny = TinyClassifier().eval()
tiny_pruned = ChannelPruner(amount=0.5).apply(tiny, target_layers=["conv"], sample_input=torch.randn(1, 3, 8, 8))
print("conv out_channels:", tiny_pruned.conv.out_channels, "| fc in_features:", tiny_pruned.fc.in_features)

import torch.nn as nn class TinyClassifier(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(3, 16, kernel_size=3, padding=1) self.pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Linear(16, 10) def forward(self, x): x = self.conv(x) x = self.pool(x) x = torch.flatten(x, 1) return self.fc(x) tiny = TinyClassifier().eval() tiny_pruned = ChannelPruner(amount=0.5).apply(tiny, target_layers=["conv"], sample_input=torch.randn(1, 3, 8, 8)) print("conv out_channels:", tiny_pruned.conv.out_channels, "| fc in_features:", tiny_pruned.fc.in_features)

conv out_channels: 8 | fc in_features: 8

Quantization¶

Quantization reduces the numerical precision used to store and compute a model's weights and activations, typically from 32-bit floats down to 8-bit integers. Instead of removing parameters, it makes every parameter cheaper to store and to multiply. The standard scheme is affine (asymmetric) quantization: a real value $x$ is mapped to an integer $q$ via a per-tensor (or per-channel) scale and zero-point,

$$q = \text{round}\left(\frac{x}{s}\right) + z, \qquad x \approx s \cdot (q - z)$$

where $s$ (scale) and $z$ (zero-point) are chosen so the tensor's actual value range maps as nearly as possible into the 8-bit range $[-128, 127]$ (or $[0, 255]$). shrinkai.compression.quantization.Quantizer supports three ways of getting there:

• Dynamic PTQ (Post-Training Quantization): weights are converted to int8 ahead of time, activations are quantized on the fly, per batch, at inference time. No calibration data needed. Only converts nn.Linear layers, Conv2d layers remain the same.
• Static PTQ: both weights and activations are converted to int8 ahead of time. This requires running a calibration pass over representative data first, so the quantizer can observe the actual range each activation takes, but in exchange, it also quantizes Conv2d layers, not just Linear.
• QAT (Quantization-Aware Training): inserts "fake quantization" nodes that simulate int8 rounding during a normal float32 training loop, so the model learns to be robust to the quantization noise before it is actually converted at the end.

Note on hardware: quantized kernels are backend-specific, fbgemm on x86 CPUs, qnnpack on ARM/mobile. torch.backends.quantized.engine must be set explicitly before quantizing (there usually is no useful default), which is itself a reminder that int8 inference is fundamentally an edge/CPU-oriented technique, not something that benefits a GPU training pipeline so far.

Note on torch.ao.quantization: this is the API Quantizer is built on, and PyTorch has marked it deprecated in favor of torchao. It still works correctly (that is what this whole section demonstrates), but expect a DeprecationWarning, migrating once torchao's API has stabilized is a known next step for shrinkai.

import torch.backends.quantized

torch.backends.quantized.engine = "qnnpack"  # use "fbgemm" instead on x86/server CPUs

import torch.backends.quantized torch.backends.quantized.engine = "qnnpack" # use "fbgemm" instead on x86/server CPUs

Why a different model for this section. resnet20's residual connections (out += identity) are a problem for quantization, for a very different reason than for ChannelPruner: PyTorch's eager-mode quantization converts a model region by region, and a raw += between two tensors sitting right at the boundary between "already quantized" and "not yet quantized" territory has no quantized kernel to fall back to. This is a general limitation of eager-mode torch.ao.quantization on any model with an addition-based skip connection, not something specific to shrinkai. We switch to vgg19_bn, which is a plain stack of Conv2d → BatchNorm → ReLU blocks with no such branch, and works cleanly end to end.

vgg = torch.hub.load(
    "chenyaofo/pytorch-cifar-models",
    "cifar10_vgg19_bn",
    pretrained=True,
    verbose=False,
).eval()

print("Parameters:", count_parameters(vgg)["total_params"])
print("Disk size (MB):", estimate_model_size_mb(vgg))

vgg = torch.hub.load( "chenyaofo/pytorch-cifar-models", "cifar10_vgg19_bn", pretrained=True, verbose=False, ).eval() print("Parameters:", count_parameters(vgg)["total_params"]) print("Disk size (MB):", estimate_model_size_mb(vgg))

Parameters: 20565834
Disk size (MB): 78.53

Dynamic PTQ¶

from shrinkai.compression.quantization import QuantConfig, Quantizer

from shrinkai.compression.quantization import QuantConfig, Quantizer

config = QuantConfig(strategy="ptq")  # dynamic (no calibrate_data)
quantizer = Quantizer(config)

dynamic_quantized = quantizer.apply(copy.deepcopy(vgg))
report = quantizer.benchmark(
    vgg,
    dynamic_quantized,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

config = QuantConfig(strategy="ptq") # dynamic (no calibrate_data) quantizer = Quantizer(config) dynamic_quantized = quantizer.apply(copy.deepcopy(vgg)) report = quantizer.benchmark( vgg, dynamic_quantized, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

/Users/elouan/Documents/projet_python/shrinkai/.venv/lib/python3.13/site-packages/torch/ao/nn/quantized/modules/utils.py:72: UserWarning: torch.quantize_per_tensor, torch.quantize_per_channel and other quantized tensor creation functions that produce tensors with dtype torch.quint8, torch.qint8, and torch.qint32 are deprecated and will be removed in a future PyTorch release. Please see https://github.com/pytorch/pytorch/issues/184982 for more information. (Triggered internally at /Users/runner/work/pytorch/pytorch/aten/src/ATen/quantized/Quantizer.cpp:116.)
  qweight = torch.quantize_per_tensor(
[W913 11:41:29.369495000 qlinear_dynamic.cpp:251] Warning: Currently, qnnpack incorrectly ignores reduce_range when it is set to true; this may change in a future release. (function operator())

                                       Benchmark Report                                       
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (FP32)) ┃ Student (Quantized) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                   20.57 M │             20.04 M │                -2.6% │
│ Model Size (Disk) │                  78.53 MB │            77.02 MB │ -1.9% (1.0x smaller) │
│ Latency / Sample  │                   3.06 ms │             3.19 ms │          1.0x faster │
│ Throughput (FPS)  │               326.5 img/s │         313.0 img/s │    +1.0x (313.0 FPS) │
│ Accuracy          │                    93.40% │              93.40% │      100.0% retained │
└───────────────────┴───────────────────────────┴─────────────────────┴──────────────────────┘

vgg19_bn's 20M parameters live almost entirely in its Conv2d layers (16 of them) plus a 3-layer Linear classifier head; dynamic PTQ only touches the latter. This is the expected, honest result of "dynamic PTQ only converts nn.Linear" applied to a convolution-based architecture, it's a legitimate, zero-risk quick win for MLP/RNN/Transformer-heavy models, but not the right tool for a CNN backbone like this one.

Static PTQ¶

Static PTQ needs a calibration dataset, a handful of representative batches used only to observe the range each activation takes, not to train anything. Passing calibrate_data to QuantConfig is what switches Quantizer from dynamic to static PTQ:

calibration_loader = DataLoader(Subset(datasets.CIFAR10(
    root="./data", train=True, download=True,
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ]),
), range(64)), batch_size=32)

config = QuantConfig(strategy="ptq", calibrate_data=calibration_loader)
quantizer = Quantizer(config)

static_quantized = quantizer.apply(copy.deepcopy(vgg))
report = quantizer.benchmark(
    vgg,
    static_quantized,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

calibration_loader = DataLoader(Subset(datasets.CIFAR10( root="./data", train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]), ), range(64)), batch_size=32) config = QuantConfig(strategy="ptq", calibrate_data=calibration_loader) quantizer = Quantizer(config) static_quantized = quantizer.apply(copy.deepcopy(vgg)) report = quantizer.benchmark( vgg, static_quantized, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

                                       Benchmark Report                                        
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (FP32)) ┃ Student (Quantized) ┃    Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                   20.57 M │              0.01 M │                -99.9% │
│ Model Size (Disk) │                  78.53 MB │            19.76 MB │ -74.8% (4.0x smaller) │
│ Latency / Sample  │                   3.11 ms │             2.47 ms │           1.3x faster │
│ Throughput (FPS)  │               322.0 img/s │         404.5 img/s │     +1.3x (404.5 FPS) │
│ Accuracy          │                    93.40% │              93.40% │       100.0% retained │
└───────────────────┴───────────────────────────┴─────────────────────┴───────────────────────┘

This time the model is genuinely ~4x smaller and faster, for almost no accuracy loss because the Conv2d backbone is now int8 too, not just the classifier head. This is the real payoff PTQ can offer for CNNs, at the cost of needing representative calibration data.

However, the number of parameters plummeted. This drop is misleading: once converted, a quantized Conv2d/Linear's weights are stored as an opaque packed buffer (_packed_params), not as an nn.Parameter anymore, so count_parameters() (which sums model.parameters()) simply stops seeing most of them. The "Model Size" row (based on the actual serialized state_dict()) is the trustworthy number here, and it agrees with the expected ~4x from dropping 32-bit floats to 8-bit integers.

Quantization-Aware Training (QAT)¶

QAT is the right tool when PTQ's accuracy loss (dynamic or static) is too large for the task at hand: instead of quantizing a finished model and hoping it survives, it inserts fake-quantize nodes that simulate int8 rounding during training, so the weights themselves adjust to deal with quantization noise.

This model needs an explicit QuantWrapper. Quantizer's QAT path (strategy="qat") prepares the model for fake-quantized training directly, but does not insert the entry/exit QuantStub/DeQuantStub boundaries that _apply_static_ptq adds automatically. Those are mandatory to actually run inference once the model is finalized. Until shrinkai handles this automatically for QAT too, wrap the model yourself with torch.ao.quantization.QuantWrapper first (exactly what static PTQ does internally).

wrapped_vgg = torch.ao.quantization.QuantWrapper(copy.deepcopy(vgg))

config = QuantConfig(strategy="qat")
quantizer = Quantizer(config)
qat_model = quantizer.apply(wrapped_vgg)

wrapped_vgg = torch.ao.quantization.QuantWrapper(copy.deepcopy(vgg)) config = QuantConfig(strategy="qat") quantizer = Quantizer(config) qat_model = quantizer.apply(wrapped_vgg)

qat_model behaves like a normal, fully differentiable float32 model, the fake-quantize nodes only simulate rounding in the forward pass, so it can be trained exactly like any other model, including through a Distiller. There is no separate teacher here: we just want to fine-tune qat_model on real labels, so we route it through HintonLoss with alpha=0.0, a convenient way to get a plain supervised cross-entropy loss out of the same Distiller API used everywhere else in this library (the "teacher" argument is still required by the API, but its output ends up weighted by zero).

from shrinkai.distillation import Distiller
from shrinkai.distillation.losses import HintonLoss

distiller = Distiller(
    teacher=copy.deepcopy(vgg),
    student=qat_model,
    criterion=HintonLoss(alpha=0.0),
    optimizer="adamw",
    lr=1e-4,
    device="cpu",
)

_ = distiller.fit(train_loader, epochs=1)

from shrinkai.distillation import Distiller from shrinkai.distillation.losses import HintonLoss distiller = Distiller( teacher=copy.deepcopy(vgg), student=qat_model, criterion=HintonLoss(alpha=0.0), optimizer="adamw", lr=1e-4, device="cpu", ) _ = distiller.fit(train_loader, epochs=1)

Epoch [01/01] Train Loss: 0.0369 - Train Acc: 99.60%

qat_model.eval()
finalized_qat = Quantizer.finalize_qat(qat_model)

report = Quantizer(config).benchmark(
    vgg,
    finalized_qat,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

qat_model.eval() finalized_qat = Quantizer.finalize_qat(qat_model) report = Quantizer(config).benchmark( vgg, finalized_qat, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

                                       Benchmark Report                                        
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (FP32)) ┃ Student (Quantized) ┃    Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                   20.57 M │              0.01 M │                -99.9% │
│ Model Size (Disk) │                  78.53 MB │            19.76 MB │ -74.8% (4.0x smaller) │
│ Latency / Sample  │                   3.20 ms │             2.63 ms │           1.2x faster │
│ Throughput (FPS)  │               312.1 img/s │         380.5 img/s │     +1.2x (380.5 FPS) │
│ Accuracy          │                    93.40% │              94.20% │       100.9% retained │
└───────────────────┴───────────────────────────┴─────────────────────┴───────────────────────┘

Even a single epoch on a small subset is enough for the fake-quantize nodes to adapt the weights, landing close to (sometimes above) static PTQ's accuracy, with the same ~4x size reduction. In practice, QAT is worth the extra training step precisely when static PTQ's calibration-only approach loses too much accuracy for the task.

Combining pruning and quantization¶

The two families of techniques act on different axes (which parameters exist vs. what precision they use) and compose freely. You actually can quantize an already pruned model. Reusing aggressively_pruned from the Pruning section:

config = QuantConfig(strategy="ptq")
quantizer = Quantizer(config)

pruned_and_quantized = quantizer.apply(copy.deepcopy(aggressively_pruned))
report = quantizer.benchmark(
    resnet,
    pruned_and_quantized,
    sample_input=sample_input,
    val_dataloader=val_loader,
    device="cpu",
)
report.show()

config = QuantConfig(strategy="ptq") quantizer = Quantizer(config) pruned_and_quantized = quantizer.apply(copy.deepcopy(aggressively_pruned)) report = quantizer.benchmark( resnet, pruned_and_quantized, sample_input=sample_input, val_dataloader=val_loader, device="cpu", ) report.show()

                                       Benchmark Report                                       
┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric            ┃ Teacher (Original (FP32)) ┃ Student (Quantized) ┃   Gain / Compression ┃
┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ Parameters        │                    0.27 M │              0.27 M │                -2.1% │
│ Model Size (Disk) │                   1.08 MB │             1.06 MB │ -1.9% (1.0x smaller) │
│ Latency / Sample  │                   0.96 ms │             0.94 ms │          1.0x faster │
│ Throughput (FPS)  │              1043.1 img/s │        1060.9 img/s │   +1.0x (1060.9 FPS) │
│ Accuracy          │                    91.80% │              83.60% │       91.1% retained │
└───────────────────┴───────────────────────────┴─────────────────────┴──────────────────────┘

On resnet20 specifically, dynamic PTQ adds little on top of ChannelPruner (the same Linear limitation from earlier: this network has a single, tiny fc layer), but the combination is exactly what you would reach for on a model with a larger classifier head, or by switching to static PTQ for the convolutional layers too.

Summary¶

Technique	Physically smaller?	Faster?	Needs data?	Needs training?	Best for
Pruner (unstructured)	No (masked)	No	No	No	Preparing a model for pruning-aware further training; sparse-hardware deployment
Pruner (structured)	No (masked)	No	No	No	Same, when you specifically need whole channels zeroed (e.g. as a step before manual pruning)
ChannelPruner	Yes	Yes	No (needs sample_input only)	Recommended after pruning, to recover accuracy	Simple, non-branching CNN/MLP backbones; real size/latency/FLOPs budget cuts
Dynamic PTQ	Yes, for Linear-heavy models	Sometimes	No	No	Quick, risk-free win on RNN/Transformer/MLP-heavy models
Static PTQ	Yes (weights + activations)	Yes	Yes (calibration only)	No	CNNs and other Conv2d-heavy models, when calibration data is available
QAT	Yes	Yes	Yes (training data)	Yes	When PTQ's accuracy loss is too large to accept

None of these are mutually exclusive with shrinkai.distillation, the most robust deployment pipeline is usually: distill (if you need a smaller architecture), prune (if some capacity is still redundant), quantize (to cut precision), and fine-tune between steps whenever accuracy needs recovering.

References¶

1. LeCun, Y., Denker, J. S., & Solla, S. A. (1989). Optimal Brain Damage. NeurIPS.
   https://proceedings.neurips.cc/paper/1989/hash/6c9882bbac1c7093bd25041881277658-Abstract.html

2. Han, S., Pool, J., Tran, J., & Dally, W. (2015). Learning both Weights and Connections for Efficient Neural Networks. NeurIPS. arXiv:1506.02626.
   https://arxiv.org/abs/1506.02626

3. Li, H., Kadav, A., Durdanovic, I., Samet, H., & Graf, H. P. (2017). Pruning Filters for Efficient ConvNets. ICLR. arXiv:1608.08710.
   https://arxiv.org/abs/1608.08710

4. Han, S., Mao, H., & Dally, W. J. (2016). Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR. arXiv:1510.00149.
   https://arxiv.org/abs/1510.00149

5. Jacob, B., Kligys, S., Chen, B., et al. (2018). Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. CVPR. arXiv:1712.05877.
   https://arxiv.org/abs/1712.05877

6. Krizhevsky, A. (2009). Learning Multiple Layers of Features from Tiny Images. Technical Report, University of Toronto.