Package: fastsparse
FastSparseCustomizable Fastai+PyTorch implementation of sparse model training methods (SET, SNFS, RigL).Warning: this repo is undergoing active developmentGetting StartedInstallpip install fastsparseSparse AlgorithmsThis network implements the following sparse algorithms:Abbr.Sparse Algorithmin FastSparseNotesstatic sparsity baselineomitDynamicSparseTrainingCallbackSETSparse Evolutionary Training(Jan 2019)DynamicSparseTrainingCallback(**SET_presets)SNFSSparse Networks From Scratch(Jul 2019)DynamicSparseTrainingCallback(**SNFS_presets)redistribution not implemented*RigLRigged Lottery(Nov 2019)DynamicSparseTrainingCallback(**RigL_presets)*Authors of the RigL paper demonstrate that using SNFS + Erdos-Renyi-Kernel distribution - redistribution outperforms SNFS + uniform sparsity + redistribution (at least on the measured benchmarks).Fastai demoWith just 4 additional lines of code, you can train your model using the latest dynamic sparse training techniques. This example achieves >99% accuracy on MNIST using a ResNet34 with only 1% of the weights.# (0) install the library# ! pip install fastsparsefromfastai.vision.allimport*# (1) import this packageimportfastsparseassparsepath=untar_data(URLs.MNIST)dls=ImageDataLoaders.from_folder(path,'training','testing')learn=cnn_learner(dls,resnet34,metrics=error_rate,pretrained=False)# (2) sparsify initial model + enforce maskssparse_hooks=sparse.sparsify_model(learn.model,model_sparsity=0.99,sparse_f=sparse.erdos_renyi_sparsity)# (3) schedule dynamic mask updatescbs=[sparse.DynamicSparseTrainingCallback(**sparse.SNFS_presets,batches_per_update=32)]learn.fit_one_cycle(5,cbs=cbs)# (4) remove hooks that enforce maskssparse_hooks.remove()Simply omit theDynamicSparseTrainingCallbackto train a fixed-sparsity model as a baseline.PyTorch demo (not implemented yet)importtorchfromtorchvisionimportmodelsdata=...model=...opt=...opt=DynamicSparseTrainingOptimizerWrapper(model,opt,**RigL_kwargs)### Modified training step# sparse_opt.step(...) will determine whether to:#  (A) take a regular opt step, or#  (B) update network connectivitydefsparse_train_step(model,xb,yb,loss_func,sparse_opt,step,pct_train):preds=model(xb)loss=loss_func(preds,yb)loss.backward()sparse_opt.step(step,pct_train)sparse_opt.zero_grad()Save/Reload demoHere is an example of saving a model and reloading it to resume training.fromfastai.vision.allimport*fromfastsparseimport*path=untar_data(URLs.MNIST_TINY)dls=ImageDataLoaders.from_folder(untar_data(URLs.MNIST_TINY))learn=cnn_learner(dls,resnet18,metrics=accuracy,pretrained=False)sparse_hooks=sparsify_model(learn.model,model_sparsity=0.9,sparse_f=erdos_renyi_sparsity)dst_kwargs={**SNFS_presets,**{'batches_per_update':8}}cbs=DynamicSparseTrainingCallback(**dst_kwargs)learn.fit_flat_cos(5,cbs=cbs)# (0) save model as usual (masks are stored automatically)save_model('sparse_tiny_mnist',learn.model,learn.opt)epochtrain_lossvalid_lossaccuracytime00.4827980.6989340.50500700:0210.3024420.6562830.51216000:0120.2386230.1756930.93562200:0230.2039080.0286190.99284700:0240.1621430.0339450.98998600:02### manually restart notebook #### (1) then recreate learner as usualfromfastai.vision.allimport*fromfastsparseimport*path=untar_data(URLs.MNIST_TINY)dls=ImageDataLoaders.from_folder(untar_data(URLs.MNIST_TINY))learn=cnn_learner(dls,resnet18,metrics=accuracy,pretrained=False)# (2) re-sparsify model (this adds the masks to the parameters)sparse_hooks=sparsify_model(learn.model,model_sparsity=0.9,sparse_f=erdos_renyi_sparsity)# <-- initial sparsity + enforce masks# (3) load model as usualload_model('sparse_tiny_mnist',learn.model,learn.opt)# (5) check validation loss & accuracy to verify we've loaded it successfullyval_loss,val_acc=learn.validate()print(f'validation loss:{val_loss}, validation accuracy:{val_acc}')# (4) optionally, continue training; otherwise remove sparsity-preserving hookssparse_hooks.remove()/home/dc/anaconda3/envs/fastai/lib/python3.8/site-packages/fastai/learner.py:53: UserWarning: Could not load the optimizer state.
  if with_opt: warn("Could not load the optimizer state.")






validation loss: 0.033944640308618546, validation accuracy: 0.9899857044219971Training with Large Batch SizesAuthors of the Rigged Lottery paper hypothesize that the effectiveness of using the gradient magnitude for determining which connections to grow is partly due to their large batch size (4096 for ImageNet). Those without access to multi-gpu clusters can achieve effective batch sizes of this size by using fastai'sGradientAccumulationcallback, which has been tested to be compatible with this package'sDynamicSparseTrainingCallback.Training with Small # of EpochsDynamic sparse training algorithms work by modifying the network connectivity during training, dropping some weights and allowing others to regrow. By default, network connectivity is modified at the end of each epoch. When training with few epochs, however, there will be few chances to explore which weights to connect. To update more frequently, inDynamicSparseTrainingCallback, setbatches_per_updateto a smaller # of batches than occur in one training epoch. Varying the number of batches per update trades off the frequency of updates with stability in making good updates.CustomizationThere are many ways to implement and test your own dynamic sparse algorithms using FastSparse.Custom Initial Sparsity Distribution:Define your own initial sparsity distribution by settingsparsify_methodinsparsify_modelto a custom function. For example, this function (included in library) will keep the first layer dense and set the remaining layers to a fixed sparsity.deffirst_layer_dense_uniform(params:list,model_sparsity:float):sparsities=[1.]+[model_sparsity]*(len(params)-1)returnsparsitiesCustom Drop CriterionWhile published papers like SNFS and RigL refer to 'drop criterion', this library implements the reverse, a 'keep criterion'. This is a function that returns a score for each weight, where the largestMscores will be andMis determined by the decay schedule. For example, both Sparse Networks From Scratch and Rigged Lottery both use the magnitude of the weights (in FastSparse:weight_magnitude).This can easily be customized in FastSparse by defining your own keep score function:defcustom_keep_scoring_function(param,opt):score=...assertparam.shape==score.shapereturnscoreThen pass your custom function into the sparse training callback:DynamicSparseTrainingCallback(...,keep_score_f=custom_keep_scoring_function)Custom Grow CriterionThe grow criterion is a function that returns a score for each weight, where the largestNscores will be andNis determined by the decay schedule. For example, Sparse Networks From Scrath grows weights according to the momentum of the gradient, while Rigged Lottery uses the magnitude of the gradient (in FastSparse,gradient_momentumandgradient_magnituderespectively).defcustom_grow_scoring_function(param,opt):score=...assertparam.shape==score.shapereturnscoreThen pass your custom function into the sparse training callback:DynamicSparseTrainingCallback(...,grow_score_f=custom_grow_scoring_function)Replication ResultsIn machine learning, is very easy for seemingly insignificant differences in algorithmic implementation to have a noticeable impact on final results. Therefore, this section compares results from this implementation to results reported in published papers.TODO...Under-The-Hood DetailsHere's what's going on.When you runsparsify_model(learn.model, 0.9), this adds sparse masks and add pre_forward hooks to enforce masks on weights during forward pass.By default, a uniform sparsity distribution is used. Change the sparsity distribution to Erdos-Renyi withsparsify_model(learn.model, 0.9, sparse_init_f=erdos_renyi), or pass in your custom function (seeCustomizationTo avoid adding pre_forward hooks, usesparsify_model(learn.model, 0.9, enforce_masks=False).When you add theDynamicSparseTrainingCallbackcallback, ... TODO complete section