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

Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.16.0             Please see https://github.com/pytorch/ao/issues/2919 for more info

source

PruneCallback


def PruneCallback(
    pruning_ratio, # Ratio of params to remove: float/int (0-1 or 0-100), or dict[layer_name, ratio] for per-layer targets (requires context='local')
    schedule, # When to prune, from `fasterai.core.schedule` (e.g. one_shot, agp)
    context, # 'local' (per-layer) or 'global' (across the whole model); per-layer dict requires 'local'
    criteria, # How to select filters to prune, from `fasterai.core.criteria`
    args:VAR_POSITIONAL, kwargs:VAR_KEYWORD
):

Basic class handling tweaks of the training loop by changing a Learner in various events

markdownParameters:

pruning_ratio: Target ratio of parameters to remove. Either a single value (0-1 or 0-100, values >1 treated as percentages) applied everywhere, or a dict mapping layer names to per-layer ratios (e.g. {'layer1': 30, 'layer3': 60}) for non-uniform pruning.
schedule: When to prune (from fasterai.core.schedule). Controls how pruning progresses over training.
context: 'local' (per-layer pruning) or 'global' (across entire model). Per-layer dict targets require 'local' — global context compares importance across layers, which is incompatible with non-uniform targets.
criteria: How to select what to prune (from fasterai.core.criteria).

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=30,        # 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': 30, 'layer3': 60},
    schedule=one_shot,
    context='local',
    criteria=large_final
)
learn.fit(10, cbs=[cb])

Overview

PruneCallback

Usage Example

See Also