Regularize Callback

Perform Group Regularization in fastai Callback system

Overview

The RegularizeCallback applies structured regularization during training to encourage weight sparsity at various granularities. This is useful as a pre-pruning step: with norm=2, each group of the chosen granularity is penalized as a unit (group lasso), within its layer. In a network where each convolution is followed by BatchNorm, the BatchNorm rescales each filter’s output, so penalizing the convolution weights does not select filters; to prepare structured pruning there, penalize the BatchNorm scale or use GroupRegularizeCallback, whose group penalty includes it.

Key Features: - Supports every granularity defined in fasterai.core.granularity.Granularities (for Conv2d: 'weight', 'shared_weight', 'channel', 'column', 'row', 'kernel', 'filter', and their shared_/slice variants) - Compatible with any criteria from fasterai.core.criteria - Optional scheduling to vary regularization strength over training

source

RegularizeCallback

def RegularizeCallback(
    criteria:Criteria | list[Criteria], # Importance criteria, e.g. large_final
    granularity:str | list[str], weight:float=0.01,
    layer_types:Type | list[Type]=Conv2d, # Module types to regularize
    schedule:Schedule | None=None, # Optional schedule for the weight
    verbose:bool=False, # Report the weight after each epoch
    *, norm:int=1, # 1: elementwise L1, 2: group lasso over `granularity`
):

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

Parameters: - granularity: Level at which to group weights, from Granularities (e.g. 'weight', 'channel', 'kernel', 'filter', 'layer') - weight: Regularization coefficient (higher = stronger regularization) - norm: 1 (default) is an elementwise L1, the same value whatever the granularity; 2 sums the L2 norm of each group of the granularity (group lasso), so each group is pushed to zero as a unit. Groups are taken within a single layer: they do not follow cross-layer couplings such as the next layer’s input channels or BatchNorm parameters.

Usage Example

Apply filter-level group lasso (norm=2), which penalizes each filter as a unit:

from fasterai.regularize.regularize_callback import RegularizeCallback
from fasterai.core.criteria import large_final

# Group lasso at filter granularity: each filter's L2 norm is penalized as a unit
cb = RegularizeCallback(
    criteria=large_final,
    granularity='filter',
    weight=0.01,
    norm=2,
    verbose=True
)

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

Typical Workflow: 1. Train with RegularizeCallback to penalize groups of weights (in convolution + BatchNorm blocks, see the caveat above) 2. After training, use PruneCallback or Pruner to remove the least important structures 3. Fine-tune the pruned model to recover any lost accuracy

Overview

RegularizeCallback

Usage Example

See Also