Model Pruning — Removing Redundant Weights for Faster Inference

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
import numpy as np

model = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(),
    nn.Linear(256, 128), nn.ReLU(),
    nn.Linear(128, 10)
)

# Train model first...
# Then prune

# ── Method 1: Unstructured magnitude pruning (zeroes individual weights) ──
# Prune 40% of weights in each Linear layer with smallest absolute magnitude
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        prune.l1_unstructured(module, name='weight', amount=0.4)
        # Creates a weight_mask: 0 for pruned, 1 for kept

# Check sparsity
total, zeros = 0, 0
for name, param in model.named_parameters():
    if 'weight' in name:
        total += param.numel()
        zeros += (param == 0).sum().item()
print(f"Global sparsity: {100 * zeros / total:.1f}%")   # ~40%

# Make pruning permanent (removes masks, keeps zero weights)
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        prune.remove(module, 'weight')

# ── Method 2: Global magnitude pruning (prune globally across all layers) ──
parameters_to_prune = [
    (module, 'weight') for _, module in model.named_modules()
    if isinstance(module, nn.Linear)
]
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.5,    # Remove 50% of weights globally (not per-layer)
)

# ── Method 3: Structured pruning (removes entire neurons — hardware friendly) ──
# Prune 30% of channels in each layer based on L2 norm of weight rows
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        prune.ln_structured(module, name='weight', amount=0.3, n=2, dim=0)
        # dim=0: prune output neurons (rows). Each pruned row = dead neuron.

# ── Method 4: Iterative magnitude pruning (gradual removal + fine-tuning) ──
def iterative_prune(model, target_sparsity, n_iterations, fine_tune_steps):
    """Gradually prune model, fine-tuning after each round."""
    amount_per_round = 1 - (1 - target_sparsity) ** (1 / n_iterations)
    for iteration in range(n_iterations):
        print(f"Pruning round {iteration+1}/{n_iterations}")
        for module in model.modules():
            if isinstance(module, nn.Linear):
                prune.l1_unstructured(module, 'weight', amount=amount_per_round)
        # Fine-tune for fine_tune_steps steps here...
        print(f"  Fine-tuning for {fine_tune_steps} steps...")
    return model

# ── LLM-specific: Attention head pruning ──
# Many attention heads are redundant after training
# Measure importance by how much model performance degrades when each head is zeroed
def get_head_importance(model, eval_loader):
    """Approximate head importance by gradient magnitude."""
    head_importances = {}
    for name, module in model.named_modules():
        if hasattr(module, 'num_heads'):
            # Importance = mean absolute value of attention weight gradients
            grads = module.weight.grad.abs().mean(dim=[0,2,3])   # Per-head
            head_importances[name] = grads.detach()
    return head_importances