Sparsify Callback

Use the sparsifier in fastai Callback system

SparsifyCallback

 SparsifyCallback (sparsity, granularity, context, criteria, schedule,
                   lth=False, rewind_epoch=0, reset_end=False,
                   save_tickets=False, model=None, round_to=None,
                   layer_type=<class 'torch.nn.modules.conv.Conv2d'>)

Sparsify model during training

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 optionnal 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)`