Pruning a VGG on CIFAR-10
How our solution shrinks a 35 MiB image classifier in two ways: fine-grained pruning, which zeroes individual weights, and channel pruning, which removes whole feature maps. Every figure on this page comes from the pretrained checkpoint and our verified code, run locally on an Apple M4.
Drop the smallest weights first
The importance of a weight is its magnitude. For a target sparsity s we zero round(s · #W) weights: the k-th smallest magnitude becomes the threshold, and only weights strictly above it survive. Drag the slider over the notebook's own 5×5 test tensor.
num_zeros = round(num_elements * sparsity)
importance = torch.abs(tensor)
threshold = importance.view(-1).kthvalue(num_zeros).values
mask = torch.gt(importance, threshold)
tensor.mul_(mask)Q2 check. At 0.75, round(18.75) = 19 weights go, the threshold is the 19th smallest magnitude (0.39), and the 6 surviving weights match the notebook's expected mask exactly.
Q3 answer. To keep exactly 10 of 25 weights we need 15 zeros, so target_sparsity = 0.6. Verified: exactly 10 nonzeros remain.
The outlined cell is the threshold weight itself. Because the mask uses strict >, it is pruned along with everything smaller.
Most weights already sit near zero
Every layer's weights form a narrow bell centered on zero with heavy tails, and the spread shrinks sharply with depth. The curves below use each layer's measured standard deviation; the bar under each one shows the measured share of weights with |w| < 0.02. The deep layers that hold most parameters are the easiest to prune.
Prune one layer at a time, measure the damage
The scan prunes a single layer at sparsities 0.4 to 0.9, evaluates on the 10,000 test images, restores the layer, and moves on: 9 layers × 6 sparsities = 54 full evaluations. These are our measured curves. The dashed line is the notebook's lower bound, 89.42%, which is 96% of the dense accuracy.
Most sensitive: conv0. It already drops at 0.4, falls under the bound at 0.6 (87.55%) and collapses to 31.81% at 0.9, because its 1,728 weights read the raw RGB image with almost no redundancy. conv7 and the classifier barely notice even 90% sparsity.
Spend the pruning where the parameters are
conv4 to conv7 hold 89.5% of all weights, so they carry the compression. The small, sensitive early layers and the classifier are pruned lightly. Switch between our chosen dictionary and the backups we pre-sized in case finetuned accuracy lands under 92.5%. Sizes count every parameter, including BatchNorm, exactly like the notebook.
Our submission. About four points of headroom under the limit.
92.79% after finetuning, target above 92.5%Train, then re-apply the mask every step
After pruning, five epochs of SGD recover accuracy. Each optimizer step nudges every weight, including the pruned ones, so the FineGrainedPruner callback multiplies the masks back in after every iteration. Red cells are pruned weights the gradient tried to revive.
Without the callback the model would drift back to dense within a few hundred steps and the size reduction would disappear.
Pruning to 20.94% drops accuracy from 92.95% to 84.48%. One epoch of masked finetuning recovers 92.52%, and the best of five epochs reaches 92.79% (fine-grained, blue). The channel-pruned model (amber) starts from 36.81% and ends at 92.30%.
Remove whole channels, get a smaller dense model
Channel pruning keeps round(C · 0.7) output channels of each conv and slices the matching input channels of the next conv. The result needs no sparse format and runs faster on any hardware.
n_keep = int(round(channels * (1. - prune_ratio)))
prev_conv.weight.set_(prev_conv.weight[:n_keep])
next_conv.weight.set_(next_conv.weight.detach()[:, :n_keep])Every inner conv loses 30% of both its inputs and outputs, so its cost roughly halves. conv0 keeps its 3 RGB inputs and conv7 keeps all 512 outputs, so those two only drop to about 0.70.
Q8.2. Latency improves 2.0× in the final run (1.6× in an earlier one). It is at best about the 2.0× MAC cut and often falls short, because fixed per-layer overhead, memory-bound BatchNorm, ReLU and pooling, and the unpruned first input and classifier do not shrink quadratically.
| measured | original | pruned | ratio |
|---|---|---|---|
| MACs | 606.2 M | 305.4 M | 1.98× |
| Parameters | 9.23 M | 5.01 M | 1.84× |
| Notebook check | 305,388,064 MACs | exact | |
| CPU latency, batch 1 | 3.6 ms | 1.8 ms | 2.0× |
| Accuracy after finetuning | 92.95% | 92.30% | −0.65 pt |
Sort channels by importance, then keep the top 70%
Naive pruning keeps the first channels whatever they are. We instead score each input channel of the next conv by the Frobenius norm of its weights, reorder the previous conv, its BatchNorm, and the next conv's inputs with the same permutation, and only then cut. The reorder alone changes nothing: our check measured a maximum logit difference of 2.2e-5 and 100% identical predictions.
Illustration with 16 channels. Bar heights span the real range of conv1's input-channel norms after sorting, 1.90 down to 0.26.
importance = torch.norm(channel_weight)
next_conv.weight.copy_(torch.index_select(next_conv.weight.detach(), 1, sort_idx))Everything checks out
Code answers were verified twice: independently on the CPU, then by the notebook's own tests during the full GPU run.