Sparsify Callback

Use the sparsifier in fastai Callback system

Overview

The SparsifyCallback integrates weight sparsification into the fastai training loop. Unlike pruning (which removes structures), sparsification zeros out individual weights while maintaining the original network shape.

Key Features: - Gradual sparsification according to a schedule - Support for Lottery Ticket Hypothesis (LTH) training - Multiple granularity levels (weight, vector, kernel, filter) - Global or local sparsification context

source

SparsifyCallback

def SparsifyCallback(
    sparsity:float | dict[str, float], # Target sparsity, a fraction in [0, 1] (0.4 = 40%), or a per-layer dict
    granularity:str, # Pruning granularity (e.g., 'weight', 'filter')
    context:str, # 'global' or 'local'
    criteria:Criteria, # Criteria for determining weights to keep
    schedule:Schedule, lth:bool=False, # Use the Lottery Ticket Hypothesis approach
    rewind_epoch:int=0, # Epoch to rewind weights to for LTH
    reset_end:bool=False, # Reset weights after pruning
    save_tickets:bool=False, # Save pruned models as "winning tickets"
    model:nn.Module | None=None, # Model to sparsify; None = learn.model
    round_to:int | None=None, # Round pruning to multiple of this value
    nm:bool=False, # Use N:M structured sparsity
    layer_type:Type[nn.Module]=Conv2d
):

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

The most important part of our Callback happens in before_batch. There, we first compute the sparsity of our network according to our schedule and then we remove the parameters accordingly.

The SparsifyCallback requires a new argument compared to the Sparsifier. Indeed, we need to know the pruning schedule that we should follow during training in order to prune the parameters accordingly.

You can use any scheduling function already available in fastai or come up with your own ! For more information about the pruning schedules, take a look at the Schedules section.

On top of that, the SparsifyCallbackcan also take many optional arguments:

lth: whether training using the Lottery Ticket Hypothesis, i.e. reset the weights to their original value at each pruning step (more information in the Lottery Ticket Hypothesis section)
rewind_epoch: the epoch used as a reference for the Lottery Ticket Hypothesis with Rewinding (default to 0)
reset_end: whether you want to reset the weights to their original values after training (pruning masks are still applied)
save_tickets: whether to save intermediate winning tickets.
model: pass a model or a part of the model if you don’t want to apply pruning on the whole model trained.
round_to: if specified, the weights will be pruned to the closest multiple value of round_to.
layer_type: specify the type of layer that you want to apply pruning to (default to nn.Conv2d)`

Usage Example

from fasterai.sparse.sparsify_callback import SparsifyCallback
from fasterai.core.schedule import cos
from fasterai.core.criteria import large_final

# Gradually sparsify to 50% using cosine schedule
cb = SparsifyCallback(
    sparsity=0.5,
    granularity='weight',
    context='global',
    criteria=large_final,
    schedule=cos
)

learn.fit(10, cbs=[cb])

sparsity is a fraction in [0, 1]: 0.5 is 50%. A percentage such as 50 is read as 0.5 for one release, with a FutureWarning.

Per-Layer Sparsity with Dict

# Different sparsity targets for different layers
cb = SparsifyCallback(
    sparsity={'conv1': 0.3, 'layer1': 0.5, 'layer2': 0.7},
    granularity='weight',
    context='local',
    criteria=large_final,
    schedule=cos
)

With Lottery Ticket Hypothesis

# Train with LTH - rewind weights to epoch 2 values after each pruning step
cb = SparsifyCallback(
    sparsity=0.9,
    granularity='weight',
    context='global', 
    criteria=large_final,
    schedule=one_cycle,
    lth=True,
    rewind_epoch=2
)

learn.fit(20, cbs=[cb])

Overview

SparsifyCallback

Usage Example

Per-Layer Sparsity with Dict

With Lottery Ticket Hypothesis

See Also