🎓 A1Learner
🔍

PyTorch Pruning

Pruning sets the smallest weights in a layer to zero. Sparsity is the share of a layer's weights that sit at zero.

Pruning one small layer

Half of this layer's twelve weights go to zero. The nonzero count drops from twelve to six.

import torch
from torch.nn.utils import prune

torch.manual_seed(0)
layer = torch.nn.Linear(4, 3)
print(int((layer.weight != 0).sum().item()))
prune.l1_unstructured(layer, name="weight", amount=0.5)
print(int((layer.weight != 0).sum().item()))

The smallest go first

Four hand-set weights make the choice visible. The two smallest magnitudes become zero, and the two biggest stay.

import torch
from torch.nn.utils import prune

layer = torch.nn.Linear(2, 2, bias=False)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[0.01, 0.5], [-2.0, 0.2]]))
prune.l1_unstructured(layer, name="weight", amount=0.5)
print(layer.weight.detach().tolist())

Making the zeros permanent

prune.remove folds the mask into the weight and drops the extra names. The zeros stay after the removal.

import torch
from torch.nn.utils import prune

torch.manual_seed(0)
layer = torch.nn.Linear(4, 3)
prune.l1_unstructured(layer, name="weight", amount=0.5)
prune.remove(layer, "weight")
print(int((layer.weight != 0).sum().item()))
print(hasattr(layer, "weight_orig"))

Example

import torch
from torch.nn.utils import prune

torch.manual_seed(0)
centers = torch.tensor([[-1.5] * 20, [1.5] * 20])
def make(rows, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.randint(0, 2, (rows,), generator=g)
    data = centers[labels] + torch.randn(rows, 20, generator=g) * 0.6
    return data, labels

train_x, train_y = make(300, 1)
test_x, test_y = make(200, 2)

torch.manual_seed(3)
model = torch.nn.Sequential(
    torch.nn.Linear(20, 16), torch.nn.ReLU(), torch.nn.Linear(16, 2))
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
loss_fn = torch.nn.CrossEntropyLoss()
for step in range(200):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

def accuracy(m):
    with torch.no_grad():
        pred = m(test_x).argmax(dim=1)
    return (pred == test_y).float().mean().item()

before = accuracy(model)
layer = model[0]
total = layer.weight.numel()
prune.l1_unstructured(layer, name="weight", amount=0.5)
effective = layer.weight.detach()
nonzero = int((effective != 0).sum().item())
after = accuracy(model)
print(f"Accuracy before pruning: {before:.2f}")
print(f"First layer weights: {total}")
print(f"Nonzero after pruning: {nonzero}")
print(f"Sparsity: {1 - nonzero / total:.2f}")
print(f"Accuracy after pruning: {after:.2f}")

Key points

  • Mask: a same-shaped tensor of ones and zeros that selects entries
  • Effective weight: the weight a layer actually uses, with its mask applied
  • Use prune.l1_unstructured(layer, name='weight', amount=0.5) to zero half of one weight
  • L1 picks by absolute value, and the smallest weights go first
  • The prune helpers live in torch.nn.utils.prune, and the name needs its own import
  • Sparsity counts the zeros in the effective weight after pruning

In this example, a small classifier trains on two seeded blob classes with 20 features. The first layer holds 320 weights into 16 hidden units. Accuracy prints before pruning, and the run reaches 1.00 on the test batch. Half of the first layer's weights are then pruned by size. The effective weight keeps 160 nonzero entries, and sparsity prints as 0.50. Accuracy after pruning prints as 1.00. The accuracy does not move here, because these blobs are easy to separate.

Note

Zeros alone do not shrink a saved file, and smaller files need extra steps

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does amount=0.5 mean in l1_unstructured?