image.png

概述

EasyPruner是一个轻量且实用的PyTorch神经网络剪枝工具包,提供了一系列即插即用的网络结构裁剪接口。任何不具备模型压缩知识背景的工程师都可通过在PyTorch工程代码中添加几行代码实现网络模型的精准瘦身,并导出用于多种平台部署的onnx模型。可使模型保持现有精度水平的情况下,成倍提升执行效率和存储效率。

特点

透明性:不需懂模型压缩知识即可流畅使用。
灵活性:无需更换训练框架,也不用为剪枝接口重构训练或评测代码,仅在原训练框架中进行几行代码增加,即插即用。
通用性:支持所有工程场景中常用的网络结构,例如ResNet、VGGNet、Inception、MobileNet等;支持多种训练框架代码,open-mmlab系列开源框架、u版YOLO系列等。
精确性:吸取神经网络剪枝的最新研究成果,提供在公开评测集上处于SOTA水平的剪枝方法,可对网络冗余连接进行精确识别,在一些常规任务上可以实现无损压缩。
实用性:剪枝后的模型可直接导出onnx,实现在NPU、ARM、GPU、CPU等多种平台的通用部署。

创新

与其他开源模型剪枝工具项目相比,本项目做了如下创新,以使剪枝工具更加强大:

  1. 提出基于ONNX的图分析的网络拓扑排序方法,可以实现前后依赖算子的自动化识别,以兼容多样化的训练框架。
  2. 提出包含自研GradDecay稀疏剪枝方法的进阶剪枝模式,该方法针对YOLOv5目标检测算法设计,可以在不增加训练成本的情况下,最大限度保留网络原始表达能力,在一般情况下实现无损剪枝,在大剪枝率情况下依旧保持低精度损失。
  3. 提出包含自研MaskL1稀疏剪枝方法的进阶剪枝模式,该方法比目前流行的Network Slimming 方法有优的剪枝精度。

    功能介绍

    本工具包提供快捷模式、进阶模式两种剪枝模式,分别具有不同的代码修改方式,以及剪枝效果,用户可自行选择:
    快捷模式:遵循 “模型加载—>裁剪—>常规微调”三段式,易用性更强,代码修改便捷
    进阶模式-GradDecay:遵循 “模型加载—>稀疏微调—>裁剪”三段式,适合YOLOv5大剪枝率剪枝,需找到训练代码中的权重更新部分,进行修改。
    进阶模式-MaskL1:遵循“模型加载—>稀疏微调—>裁剪—>微调”四段式,通用性强,效果比快捷模式好。

    设计架构图

image.png