Package: rigl-torch
rigl-torchWarning: This repository is still in active development, results are not yet up to the rigl paper spec. Coming soon!An open source implementation of Google Research's paper (Authored byUtku Evci, an AI Resident @ Google Brain):Rigging the Lottery: Making All Tickets Winners (RigL)in PyTorch as versatile, simple, and fast as possible.You only need to add2 lines of codeto your PyTorch project to use RigL to train your model with sparsity!ImageNet ResultsResults aren't quite as complete as the original paper, howeverthis implementation outperforms the original by 1%!ArchitectureSparsity %S. DistributionTop-1OriginalTop-1ResNet5090%Uniform73%72%Other Implementations:View the TensorFlow implementation (also the original)here!Additionally, it is also implemented invanilla pythonandgraphcore.Contributions Beyond the Paper:Gradient Accumulation:Motivation:The papercites their experiments for ImageNet being done using a batch size of 4096, which isn't practical for everyone since to do so you need 32 Tesla V100s to store that many 224x224x3 images in VRAM.Following this, if you are using a significantly small batch size for training (ie. bs=1024 for ImageNet), RigL may perform suboptimally due to instantaneous gradient information being quite noisy. To remedy this, I have introduced a solution for "emulating" larger batch sizes for topology modifications.Method:In regular dense training gradients are calculated per batch essentially averaging the loss of each sample and taking the derivative w.r.t the parameters. This means that if your batch size is 1024, the gradients are the accumulated average over 1024 data samples.In normal RigL the grow/drop perturbations scoring is being done on 1 batch (everydeltabatches, typicallydelta=100) and replaces the backpropagation step for that iteration. So you can see that if the batch size is significantly small, the topology modifications are being done on a very small amount of data, thus missing some potential signal from the dataset. In dense training, this is a balancing act (too large batch sizes have diminishing returns and can harm exploration, making it more likely to fall in a local minimum).Ifgradient_accumulation_nis > 1, then when RigL wants to make a topology modification it essentially takesnot onlythe current batch's gradients, but also the previousgradient_accumulation_n - 1batch's gradients. It then averages them element-wise, and uses this new matrix to score the grow/drop perturbations.Note:gradient_accumulation_nhas to be within the interval [1,delta). Ifgradient_accumulation_n== 1, then nothing has changed from the paper's spec. Ifgradient_accumulation_n== (delta- 1), RigL will score based on every single batch from the previous RigL step to the current one.Results:Setting thegradient_accumulation_nto a value > 1 increases performance on ImageNet by about 0.3-1% when using a batch size of 1024. In order to get the best results from batch size 1024 (for ImageNet),you should also multiply thedeltavalue by 4. This is because with a batch size of 4096, you are doing 4x less RigL steps (4096/1024 = 4) than if you used a batch size of 1024.User Setup:pip install rigl-torchContributor Setup:Clone this repository:git clone https://github.com/McCrearyD/rigl-torchCd into repo:cd rigl-torchInstall dependencies:pip install -r requirements.txtInstall package (-eallows for modifications):pip install -e .Usage:Run the tests by doingcd rigl-torch, thenpytest.I have provided some examples of training scripts that wereslightlymodified to add RigL's functionality. It adds a few parser statements, and only 2 required lines of RigL code usage to work! See them with links to the originals here:ImageNet|RigL|Original|RigL + SageMakerMNIST|RigL|OriginalOR more impressively,you can use the pruning power of RigL by adding 2 lines of code to your already existing training script! Here is how:fromrigl_torch.RigLimportRigLScheduler# first, create your modelmodel=...# note: only tested on torch.hub's resnet networks (ie. resnet18 / resnet50)# create your dataset/dataloaderdataset=...dataloader=...# define your optimizer (recommended SGD w/ momentum)optimizer=...# RigL runs best when you allow RigL's topology modifications to run for 75% of the total training iterations (batches)# so, let's calculate T_end according to thisepochs=100total_iterations=len(dataloader)*epochsT_end=int(0.75*total_iterations)# ------------------------------------ REQUIRED LINE # 1 ------------------------------------# now, create the RigLScheduler objectpruner=RigLScheduler(model,# model you createdoptimizer,# optimizer (recommended = SGD w/ momentum)dense_allocation=0.1,# a float between 0 and 1 that designates how sparse you want the network to be# (0.1 dense_allocation = 90% sparse)sparsity_distribution='uniform',# distribution hyperparam within the paper, currently only supports `uniform`T_end=T_end,# T_end hyperparam within the paper (recommended = 75% * total_iterations)delta=100,# delta hyperparam within the paper (recommended = 100)alpha=0.3,# alpha hyperparam within the paper (recommended = 0.3)grad_accumulation_n=1,# new hyperparam contribution (not in the paper)# for more information, see the `Contributions Beyond the Paper` sectionstatic_topo=False,# if True, the topology will be frozen, in other words RigL will not do it's job# (for debugging)ignore_linear_layers=False,# if True, linear layers in the network will be kept fully densestate_dict=None)# if you have checkpointing enabled for your training script, you should save# `pruner.state_dict()` and when resuming pass the loaded `state_dict` into# the pruner constructor# -------------------------------------------------------------------------------------------...morecode...forepochinrange(epochs):fordataindataloader:# do forward pass, calculate loss, etc....# instead of calling optimizer.step(), wrap it as such:# ------------------------------------ REQUIRED LINE # 2 ------------------------------------ifpruner():# -------------------------------------------------------------------------------------------# this block of code will execute according to the given hyperparameter schedule# in other words, optimizer.step() is not called after a RigL stepoptimizer.step()# it is also recommended that after every epoch you checkpoint your training progress# to do so with RigL training you should also save the pruner object state_dicttorch.save({'model':model.state_dict(),'pruner':pruner.state_dict(),'optimizer':optimizer.state_dict()},'checkpoint.pth')# at any time you can print the RigLScheduler object and it will show you the sparsity distributions, number of training steps/rigl steps, etc!print(pruner)# save modeltorch.save(model.state_dict(),'model.pth')