--- title: SparsifyCallback keywords: fastai sidebar: home_sidebar summary: "Use the sparsifier in fastai Callback system" description: "Use the sparsifier in fastai Callback system" nb_path: "nbs/02_sparsify_callback.ipynb" ---
path = untar_data(URLs.PETS)
files = get_image_files(path/"images")
def label_func(f): return f[0].isupper()
dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64))
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.
learn = cnn_learner(dls, resnet18, metrics=accuracy)
learn.unfreeze()
learn.fit_one_cycle(3)
Let's now try adding some sparsity in our model
learn = cnn_learner(dls, resnet18, metrics=accuracy)
learn.unfreeze()
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.
sp_cb = SparsifyCallback(end_sparsity=50, granularity='weight', method='local', criteria=large_final, sched_func=sched_cos)
learn.fit_one_cycle(3, cbs=sp_cb)
Surprisinlgy, our network that is composed of $50 \%$ of zeroes performs reasonnably well when compared to our plain and dense network.
The SparsifyCallback
also accepts a list of sparsities, corresponding to each layer of layer_type
to be pruned. Below, we show how to prune only the intermediate layers of ResNet-18.
learn = cnn_learner(dls, resnet18, metrics=accuracy)
learn.unfreeze()
sparsities = [0, 0, 0, 0, 0, 0, 50, 50, 50, 50, 50, 50, 50, 50, 0, 0, 0, 0, 0, 0]
sp_cb = SparsifyCallback(end_sparsity=sparsities, granularity='weight', method='local', criteria=large_final, sched_func=sched_cos)
learn.fit_one_cycle(4, cbs=sp_cb)
On top of that, the SparsifyCallback
can also take many optionnal arguments:
start_sparsity
: the sparsity that the schedule will use as a starting point (default to 0)start_epoch
: the epoch at which the schedule will start pruning (default to 0)end_epoch
: the epoch at which the schedule will stop pruning (default to the training epochs passed in fit
)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)model
: pass a model or a part of the model if you don't want to apply pruning on the whole model trained.layer_type
: specify the type of layer that you want to apply pruning to (default to nn.Conv2d)For example, we correctly pruned the convolution layers of our model, but we could imagine pruning the Linear Layers of even only the BatchNorm ones !