Prune Callback

Use the pruner in fastai Callback system

Overview

The PruneCallback integrates structured pruning into the fastai training loop. Unlike sparsification (which zeros weights), pruning physically removes network structures (filters, channels) to reduce model size and computation.

Key Differences from SparsifyCallback: - Removes structures entirely (not just zeros) - Uses torch-pruning library for dependency handling - Supports various pruning criteria and schedules

source

PruneCallback

def PruneCallback(
    pruning_ratio, # Filters to remove, a fraction in [0, 1] (0.4 = 40%), or a per-layer dict
    schedule, # When to prune, from `fasterai.core.schedule` (e.g. one_shot, agp)
    context, # 'local' or 'global'; a dict of ratios needs 'local'
    criteria, # How to select filters to prune, from `fasterai.core.criteria`
    *args, **kwargs
):

Prune the model during training, with fasterai.prune.Pruner. A prune replaces the parameters of every layer it shrinks, so the optimizer is re-pointed at the live ones after each one: the state of a replaced parameter is carried over, sliced like the parameter (momentum, second moments, step), the state of a parameter that survived the prune is kept, and so are the groups, the hypers, fastai’s no-weight-decay and force-train marks and the freeze — torch-pruning re-creates the parameters it replaces requiring grad, and a frozen group stays frozen.

Usage Example

from fasterai.prune.prune_callback import PruneCallback
from fasterai.core.schedule import agp, one_shot
from fasterai.core.criteria import large_final

# Uniform: prune 30% of parameters using automated gradual pruning
cb = PruneCallback(
    pruning_ratio=0.3,       # Remove 30% of parameters
    schedule=agp,            # Gradual pruning (cubic decay)
    context='global',        # Prune globally across all layers
    criteria=large_final     # Keep weights with largest magnitude
)
learn.fit(10, cbs=[cb])

# Per-layer: prune each layer to a different ratio (requires context='local')
cb = PruneCallback(
    pruning_ratio={'layer1': 0.3, 'layer3': 0.6},
    schedule=one_shot,
    context='local',
    criteria=large_final
)
learn.fit(10, cbs=[cb])

Overview

PruneCallback

Usage Example

See Also