非约束通道剪枝 · 论文

UPSCALE: 非约束通道剪枝 @ ICML 2023
Alvin Wan, Hanxiang Hao, Kaushik Patnaik, Yueyang Xu, Omer Hadad, David Güera, Zhile Ren, Qi Shan

通过从现有的剪枝器中去除约束，我们通过平均提高2.1个点来提高ImageNet后训练剪枝模型的准确性 - 有益于DenseNet (+16.9), EfficientNetV2 (+7.9), 和ResNet (+6.2)。此外，对于这些非约束剪枝模型，UPSCALE通过基线导出提高了2倍以上的推理速度。

快速开始

安装我们的包。

pip install apple-upscale

使用默认幅度剪枝器进行掩码和通道剪枝。

import torch, torchvision
from upscale import MaskingManager, PruningManager

x = torch.rand((1, 3, 224, 224), device='cuda')
model = torchvision.models.get_model('resnet18', pretrained=True).cuda()  # get any pytorch model
MaskingManager(model).importance().mask()
PruningManager(model).compute([x]).prune()

自定义剪枝

我们提供了一些开箱即用的剪枝启发式方法

• 幅度（《L1》和《L2》）
• LAMP
• FPGM
• HRank

您可以将所需的启发式算法传递给UpscaleManager.mask方法调用。您还可以在UpscaleManager.mask中配置剪枝比例。值为0.25表示25%的通道被设置为0。

from upscale.importance import LAMP
MaskingManager(model).importance(LAMP()).mask(amount=0.25)

您也可以使用任何合适的方法将通道置零。

model.conv0.weight[:, 24] = 0

然后，运行我们的导出。

PruningManager(model).compute([x]).prune()

高级

您可能希望直接访问网络片段以构建高度定制的剪枝算法。

for segment in MaskingManager(model).segments():
    # prune each segment in the network independently
    for layer in segment.layers:
        # layers in the segment

开发

注意：有关核心导出算法代码的组织方式的更多详细信息，请参阅src/upscale/pruning/README.md。

克隆并设置。

git clone git@github.com:apple/ml-upscale.git
cd upscale
pip install -e .

运行测试。

py.test src tests --doctest-modules

论文

按照开发安装说明操作，以便在paper/下可用论文代码。

要运行基线无约束导出，将baseline=True传递给PruningManager.prune。

PruningManager(model).compute([x]).prune(baseline=True)

要重现论文结果，请运行

python paper/main.py resnet18

将任何模型插入到torchvision.models命名空间中。

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 dataframe

引用

如果您认为这对您的研究有帮助，请考虑引用

@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}
}