Package: apple-upscale
Unconstrained Channel Pruning ·PaperUPSCALE: Unconstrained Channel Pruning@ICML 2023Alvin Wan,Hanxiang Hao,Kaushik Patnaik,Yueyang Xu,Omer Hadad,David Güera,Zhile Ren,Qi ShanBy removing constraints from existing pruners, we improve ImageNet accuracy for post-training pruned models by 2.1 points on average - benefiting DenseNet (+16.9), EfficientNetV2 (+7.9), and ResNet (+6.2). Furthermore, for these unconstrained pruned models, UPSCALE improves inference speeds by up to 2x over a baseline export.Quick StartInstall our package.pipinstallapple-upscaleMask and prune channels, using the default magnitude pruner.importtorch,torchvisionfromupscaleimportMaskingManager,PruningManagerx=torch.rand((1,3,224,224),device='cuda')model=torchvision.models.get_model('resnet18',pretrained=True).cuda()# get any pytorch modelMaskingManager(model).importance().mask()PruningManager(model).compute([x]).prune()Customize PruningWe provide a number of pruning heuristics out of the box:Magnitude (L1andL2)LAMPFPGMHRankYou can pass the desired heuristic into theUpscaleManager.maskmethod call. You can also configure the pruning ratio inUpscaleManager.mask. A value of0.25means 25% of channels are set to zero.fromupscale.importanceimportLAMPMaskingManager(model).importance(LAMP()).mask(amount=0.25)You can also zero out channels using any method you see fit.model.conv0.weight[:,24]=0Then, run our export.PruningManager(model).compute([x]).prune()AdvancedYou may want direct access to network segments to build a heavily-customized pruning algorithm.forsegmentinMaskingManager(model).segments():# prune each segment in the network independentlyforlayerinsegment.layers:# layers in the segmentDevelopmentNOTE:Seesrc/upscale/pruning/README.mdfor more details on how the core export algorithm code is organized.Clone and setup.gitclonegit@github.com:apple/ml-upscale.gitcdupscale
pipinstall-e.Run tests.py.test src tests --doctest-modulesPaperFollow the development installation instructions to have the paper code underpaper/available.To run the baseline unconstrained export, passbaseline=TruetoPruningManager.prune.PruningManager(model).compute([x]).prune(baseline=True)To reproduce the paper results, runpythonpaper/main.pyresnet18Plug in any model in thetorchvision.modelsnamespace.usage: main.py [-h] [--side {input,output} [{input,output} ...]]
               [--method {constrained,unconstrained} [{constrained,unconstrained} ...]]
               [--amount AMOUNT [AMOUNT ...]] [--epochs EPOCHS] 
               [--heuristic {l1,l2,lamp,fpgm,hrank}] [--global] [--out OUT] 
               [--force] [--latency] [--clean]
               model

positional arguments:
  model                 model to prune

options:
  -h, --help            show this help message and exit
  --side {input,output} [{input,output} ...]
                        prune which "side" -- producers, or consumers
  --method {constrained,unconstrained} [{constrained,unconstrained} ...]
                        how to handle multiple branches
  --amount AMOUNT [AMOUNT ...]
                        amounts to prune by. .6 means 60 percent pruned
  --epochs EPOCHS       number of epochs to train for
  --heuristic {l1,l2,lamp,fpgm,hrank}
                        pruning heuristic
  --global              apply heuristic globally
  --out OUT             directory to write results.csv to
  --force               force latency rerun
  --latency             measure latency locally
  --clean               clean the dataframeCitationIf you find this useful for your research, please consider citing@inproceedings{wan2023upscale,
  title={UPSCALE: Unconstrained Channel Pruning},
  author={Alvin Wan and Hanxiang Hao and Kaushik Patnaik and Yueyang Xu and Omer Hadad and David Guera and Zhile Ren and Qi Shan},
  booktitle={ICML},
  year={2023}
}