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.7898,  0.7097,  0.5723,  ...,  0.1163, -1.1751,  0.2450],
          [-2.9112,  0.8402,  0.9179,  ...,  0.0063,  1.1288, -0.5749],
          [-0.4368, -0.2378,  0.2966,  ...,  0.5499, -0.3741, -1.0524],
          ...,
          [ 0.4741,  1.6081, -0.7943,  ...,  0.6774, -2.6864,  1.5412],
          [ 1.0689,  1.9963, -1.0343,  ...,  0.1956,  0.1118,  1.8849],
          [-1.4006,  0.5151,  1.2845,  ...,  1.0530,  0.6049,  2.1309]],

         [[ 0.2921, -0.6453, -0.3414,  ...,  1.2958,  1.1668, -0.2857],
          [-0.5511,  0.5821, -0.7695,  ...,  1.2242, -0.8694,  0.8368],
          [ 0.4427,  0.8401,  1.6323,  ...,  0.5839, -0.6443, -1.8225],
          ...,
          [-0.3913, -0.4440, -1.7214,  ...,  0.4686, -0.2093, -0.5390],
          [-0.5736, -1.0965,  0.8492,  ..., -0.3786, -1.0599,  0.2434],
          [ 1.1565, -0.1106,  0.1925,  ...,  0.2077,  0.4050,  1.3106]],

         [[-0.3615, -2.0190, -0.0304,  ..., -0.1565, -1.1915,  1.9885],
          [ 0.9545, -0.4652, -0.2961,  ...,  0.1086, -0.4314, -0.0212],
          [ 1.4138,  0.2314,  0.0234,  ..., -1.3838,  0.6656,  0.0602],
          ...,
          [ 0.3877,  1.4464, -0.9274,  ...,  1.2363,  0.2762,  1.5912],
          [-1.9060,  0.8510, -0.7314,  ...,  0.4636,  1.8307,  0.4928],
          [ 1.7752, -0.1500,  0.3221,  ..., -0.0986,  0.1571,  1.0670]]]]), # Input used to trace layer dependencies
    *args, **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.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.

See Also

  • PruneCallback - Apply structured pruning during fastai training
  • Criteria - Different importance measures for selecting what to prune
  • Schedules - Control pruning progression during training
  • Sparsifier - Unstructured pruning (zeroing weights without removing them)
  • Parametrize - Why a model whose weights are computed is refused here
  • torch-pruning documentation - The underlying library used by Pruner

Tests live in nbs/tests/test_pruner.ipynb.