Pruner

Remove useless filters to recreate a dense network

Overview

The Pruner class provides structured pruning capabilities using the torch-pruning library. Unlike unstructured pruning (which zeros individual weights), structured pruning removes entire filters/channels, resulting in a genuinely smaller and faster model.

Key Features: - Automatic dependency handling across layers - Support for both local (per-layer) and global (cross-layer) pruning - Automatic detection and handling of attention layers in transformers - Compatible with various importance criteria from fasterai.core.criteria

Sparsifier vs Pruner: When to Use Which?

Aspect Sparsifier Pruner
What it removes Individual weights (unstructured) Entire filters/channels (structured)
Model size Same architecture, sparse weights Smaller architecture
Speedup Requires sparse hardware/libraries Immediate speedup on any hardware
Accuracy impact Generally lower at same sparsity May need fine-tuning
Best for Research, sparse-aware inference Production deployment

source

Pruner

def Pruner(
    model, pruning_ratio, # Filters to remove, a fraction in [0, 1] (0.4 = 40%), or a per-layer dict
    context, # 'local' (per-layer) or 'global' (across the whole model)
    criteria, # How to select filters to prune, from `fasterai.core.criteria`
    schedule:function=linear_scheduler, # How the ratio progresses over the pruning steps
    ignored_layers:NoneType=None, # Layers to leave untouched; None = output Linear and attention qkv
    example_inputs:torch.Tensor=tensor([[[[ 0.8956,  0.4912,  1.4266,  ..., -0.4668,  0.1172,  1.4508],
          [-0.1316,  1.7119, -0.8623,  ..., -1.0894,  1.5063,  1.1540],
          [-0.6411, -0.7524,  0.4360,  ...,  0.7936, -1.0228, -0.1795],
          ...,
          [-0.5135,  1.0052, -0.9135,  ..., -0.0970, -2.6361, -1.5480],
          [ 1.9382,  0.9718, -1.7697,  ..., -1.2244, -0.4163, -0.1126],
          [ 0.3264,  1.6185,  1.5093,  ...,  1.5141, -0.8404,  0.0138]],

         [[ 0.0556, -0.3619, -1.5456,  ..., -0.4733,  1.3603,  0.0511],
          [ 0.8676,  0.1257,  0.0238,  ..., -0.5539,  1.1292, -1.6156],
          [ 0.9638, -0.1649,  0.3306,  ..., -0.4470,  0.6638,  1.2198],
          ...,
          [ 1.1000, -0.3878, -0.6090,  ..., -0.7813,  0.5290, -0.8600],
          [ 1.4097,  2.1158,  0.7775,  ..., -1.0715, -2.7072, -1.2790],
          [ 0.9531, -1.1428, -1.0392,  ...,  1.0148, -1.6663,  0.5305]],

         [[ 0.9750, -0.8788,  0.4023,  ..., -0.7696,  0.5765,  0.8145],
          [-0.1471, -0.4398, -0.3146,  ...,  0.7080, -1.3487,  1.3047],
          [ 0.5741, -0.1023, -1.2538,  ...,  0.8947,  1.1977, -0.5741],
          ...,
          [-0.8167,  0.3180,  0.3544,  ..., -0.3721,  1.0607,  0.2923],
          [-1.1232,  1.0388, -1.4038,  ...,  0.4532,  0.9947,  0.4479],
          [ 1.2073, -0.1390, -0.1286,  ...,  1.9360,  1.4794, -0.8309]]]]), # Input used to trace layer dependencies
    *args, reg:float=0.0, # Group penalty coefficient applied by `regularize`; 0 = no penalty
    alpha:float=4, # The least important channel of a group gets 2**alpha times the penalty of the most important
    **kwargs
):

Structured pruning for neural networks using torch_pruning

source

Pruner.prune_model

def prune_model():

Execute one pruning step and restore attention layer configurations

source

Pruner.regularize

def regularize(
    scale:float=1.0, # Multiplies the penalty, for gradients that are scaled (AMP loss scale, accumulation)
):

Add the group penalty to the current gradients, between backward and the optimizer step

source

Pruner.group_importance

def group_importance(
    group
):

Compute importance scores for a dependency group

source

Pruner.print_sparsity

def print_sparsity()->None:

Print pruning report showing channel counts and parameter reduction

Usage Examples

Let’s try the Pruner with a VGG16 model

model = resnet18()
pruner = Pruner(model, 0.3, 'local', large_final)   # remove 30% of the filters
pruner.prune_model()

pruning_ratio is a fraction in [0, 1]: 0.3 is 30%. A percentage such as 30 is read as 0.3 for one release, with a FutureWarning.

Per-layer targets are a dict, where a ratio of 0 leaves that layer alone; layers the dict does not name follow default_pruning_ratio (0 by default):

pruner = Pruner(model, {'layer1': 0.2, 'layer2': 0.4, 'layer3': 0}, 'local', large_final)

Layer dependencies are traced through autograd, so the model must have at least one parameter that requires grad: on a frozen checkpoint, call model.requires_grad_(True) first.

Group sparse learning

With reg > 0, regularize() adds DepGraph’s group penalty (Fang et al., CVPR 2023) to the gradients, so that training pushes whole coupled channels towards zero before they are removed:

pruner = Pruner(model, 0.3, 'local', large_final, reg=1e-4)
for xb, yb in dl:
    loss_func(model(xb), yb).backward()
    pruner.regularize()          # after backward, before the step
    opt.step(); opt.zero_grad()
pruner.prune_model()

Every channel of a group receives reg * gamma * w on each tensor it spans: the convolution’s filter, the batch-norm weight, and the input slices of the layers that read it. gamma goes from 1 for the most important channel of the group to 2**alpha for the least important, on the square root of the criteria score. Only the groups the pruner will prune are penalized (a per-layer dict leaves the others alone), biases and frozen parameters are not, a depthwise convolution’s own kernels are neither scored nor penalized, and a group whose importance is constant is skipped. The groups are rebuilt after every prune_model(). With reg=0 (the default) nothing changes. If the gradients are scaled (a mixed-precision loss scale, gradient accumulation), pass that factor as scale. In a fastai Learner, GroupRegularizeCallback makes this call at every optimizer step.

See Also

Tests live in nbs/tests/test_pruner.ipynb.