On this page
Pruning deletes the weights a trained network barely uses. Neural networks are overparameterized (they have far more weights than the task needs), so removing those weights gives a smaller, cheaper model that keeps what it learned. Two choices shape every pruning recipe: what to cut (single weights or whole filters) and how to recover the lost accuracy (fine-tune once, or prune and fine-tune in rounds).
Magnitude pruning runs once, after training, with no data: it ranks the trained weights by size and sets a keep-or-cut mask that every later forward pass (one trip of an input through the network) multiplies by until PyTorch bakes it in. Fine-tuning is then ordinary training for a few epochs (15 in the run below); iterative and gradual pruning repeat the cycle (Section 6).
The companion repo's run uses a ResNet-18 (an 18-layer residual network for images) trained on STL-10 (a dataset of 96×96 color images in 10 classes) (Figure 1).
Pruning is the first stage of Deep Compression (Han et al., 2016), built because networks were too compute and memory hungry for embedded hardware. Huffman coding, its last stage, gives frequent values shorter codes.
The compressed models fit in on-chip SRAM (the fast memory inside the chip) instead of off-chip DRAM (the larger, slower memory outside it), and ran 3× to 4× faster per layer with 3× to 7× better energy efficiency.
Unstructured pruning removes individual weights wherever they are; the matrix keeps its shape but fills with zeros. Structured pruning removes whole filters or channels (a filter is one pattern detector in a convolutional layer, a channel the output sheet it produces), so the layer can be rebuilt smaller and runs faster on ordinary hardware.
The listing does unstructured pruning by L1 magnitude (absolute size): prune_smallest zeros the smallest share amount of every convolution and linear layer's weights, in place, and the last line cuts half of model, the trained ResNet-18.
import torch.nn as nn
import torch.nn.utils.prune as prune
KINDS = (nn.Conv2d, nn.Linear)
def prune_smallest(model, amount):
for layer in model.modules():
if isinstance(layer, KINDS):
prune.l1_unstructured(layer, "weight", amount)
prune_smallest(model, amount=0.5)Each call attaches a keep-or-cut mask to the weights it touches.
Accuracy vs sparsity (a curve through the run's two measured points)
slider setting measured on ResNet-18, unstructured
Figure 5. Accuracy falls with sparsity, never below 10% chance.
PyTorch deletes nothing at first: the weight a forward pass uses is the mask line of Section 1, built from two stored tensors (Figure 6).
prune.remove.The listing reads both stored tensors of one pruned convolution, checks their product, then bakes the mask in. Its w and m are the mask line's.
conv = model.layer1[0].conv1
w = conv.weight_orig
m = conv.weight_mask
# what forward uses
w_prime = conv.weight
assert torch.equal(w_prime, w * m)
prune.remove(conv, "weight")
The companion repo's run prunes its ResNet-18 with prune_smallest at amount=0.5. The listing computes the sparsity line: it adds up the weights into N, counts the exact zeros, and divides to get s (Figure 7).
amount is the fraction of each layer's weights set to zero. Raising it costs more accuracy before fine-tuning; lowering it keeps more weights and saves less, amount = 0 changes nothing and amount = 1 zeroes every weight. PyTorch has no default (amount is a required argument of L1Unstructured).
N = cut = 0
for layer in model.modules():
if isinstance(layer, KINDS):
w = layer.weight
N += w.numel()
cut += (w == 0).sum().item()
s = cut / N # 0.5 after the cutThe fix is to let the survivors adjust. The run trains only the final classifier layer for 15 epochs (an epoch is one full pass over the training data) with SGD (stochastic gradient descent) with momentum. Batch-norm statistics also refresh during those epochs, because the whole model is in train mode, so the recovery is not the classifier's work alone.
The listing wraps that recipe as fine_tune: SGD with η = 1e-3 and μ = 0.9 on the final layer fc only, then the standard training loop. train_loader yields STL-10 batches, criterion is nn.CrossEntropyLoss() (averaged over the batch), and y holds the true labels.
def fine_tune(model, epochs=15):
params = model.fc.parameters()
opt = torch.optim.SGD(params, lr=1e-3, momentum=0.9)
# BN statistics refresh too
model.train()
for epoch in range(epochs):
for x, y in train_loader:
opt.zero_grad()
out = model(x)
loss = criterion(out, y)
loss.backward()
opt.step()The loss falls from 0.40 to 0.23 over the 15 epochs and accuracy ends at 92.74%, above the unpruned 92.58% (Figure 1).
The recipe runs on a 6 by 6 grid that stands for the final layer, in five steps with the run's real accuracy numbers.
Figure 8. Cut, measure, retrain: accuracy comes back.
One-shot pruning prunes once, then fine-tunes. Iterative pruning prunes a little, fine-tunes, and repeats, so the survivors adapt before the next cut, and goes further for the same accuracy (Figure 9).
The listing is iterative pruning in five rounds. A repeated l1_unstructured call ranks only the weights still alive, so each round cuts 20% of what is left (about 67% gone after five rounds), and fine_tune retrains in between.
for _ in range(5):
prune_smallest(model, amount=0.2)
fine_tune(model)Gradual pruning folds the same idea into training itself: the mask is updated on a schedule while the network keeps learning.
In a real project pruning is a short stage between training and export: prune, fine-tune (The prune and fine-tune loop), then bake the masks in with prune.remove, which stores the masked weight as a plain parameter again (Figure 6).
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
KINDS = (nn.Conv2d, nn.Linear)
for layer in model.modules():
if isinstance(layer, KINDS):
# bake the mask in
prune.remove(layer, "weight")
sd = model.state_dict()
torch.save(sd, "pruned.pt")weight_orig and weight_mask, and a forward pre-hook multiplies them on every forward pass, so memory for those weights roughly doubles until prune.remove runs (PyTorch pruning tutorial).prune.global_unstructured ranks all chosen layers together, so 20% global sparsity left one layer about 5% sparse and another 22% in the PyTorch tutorial's LeNet run.prune.remove has weight_orig and weight_mask keys that an unpruned model will not load.Deep Compression (Figure 2) combines the two: prune, then store the survivors in fewer bits (Post-Training Quantization and Quantization-Aware Training). Pruning starts from a trained network; Knowledge Distillation instead trains a small network from a large one, and neural architecture search looks for a small one from the start (Neural Architecture Search Basics).