模型压缩算法库 MMRazor 全面升级,更灵活、更自动

MMRazor 是 OpenMMLab 开源项目里模型压缩领域的算法工具箱,目前支持了知识蒸馏、模型通道剪枝和模型结构搜索,模型量化部分正在开发中。

量化设计 RFC(request for comment)链接也放在这啦,欢迎各位小伙伴去提需求和建议,一起共建更好用的量化工具!

https://github.com/open-mmlab/mmrazor/issues/347

2022 年 9 月,我们发布了 MMRazor 1.0,全面升级了架构设计,适配全新的 OpenMMLab 2.0 体系

新版本的主要特点如下:

  • 非侵入式,将模型与模型压缩算法的实现解耦,不需更改模型一行代码即可压缩
  • 全面性,支持 OpenMMLab 所有 repo,压缩算法可以快速应用在不同任务上
  • 灵活性,对模型压缩算法进行了解耦,可以快速组合不同算法中的不同 trick

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第1张图片

MMRazor 支持算法类型

我们接下来将从新版本特点、重点模块两个维度,对全新的 MMRazor 1.0 进行介绍。

新版本特点

全面支持 OpenMMLab

在 OpenMMLab 1.0 时期,虽然各个算法库是基于一套统一的设计模式进行开发的,但还是会有一些不统一的地方,比如 apis 和 tools,也就导致了 MMRazor 在适配各个算法库时,需要适配每个算法库的 apis 和 tools。

在 OpenMMLab 2.0 中,我们统一了各个算法库模型、数据、训练和评测等功能的接口,使 MMRazor 可以直接应用在 OpenMMLab 不同的算法库上,不需要再额外去适配每个算法库,适配前后目录结构对比如下图所示:

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第2张图片

Config 更加简洁

在 OpenMMLab 1.0 时期,虽然 Config 有继承、导入等功能,但都是局限在同一算法库内,MMRazor 如果想用其他算法库里的 Config,必须从其他算法库中 copy-paste。比如要对 resnet-18 蒸馏,必须先从 MMClassification 中拷贝 resnet 18 的一系列 Config(模型、数据、训练等)。

OpenMMLab 2.0 新增了基础库 MMEngine,通过 MMEngine,可以实现 Config 的跨库调用,只需要在 Config 路径前加上目标 Config 对应的算法库包名, 再也不用从其他算法库拷贝大量的 Config 了,跨库调用的流程如下图所示:

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第3张图片

知识蒸馏全新升级

相较于知识蒸馏算法本身,知识蒸馏的算法实现会更加繁琐一些。很多知识蒸馏算法都是在围绕如何设计知识蒸馏的损失函数,在实现过程中,为了获取到要计算蒸馏损失的特征图,往往要通过改写模型 forward 的方式。

在 OpenMMLab 1.0 时期,MMRazor 支持了特征图的可配置获取,但只局限于 nn.Module 的输出,无法获得某个具体的函数的输入输出。

此次适配 OpenMMLab 2.0 的过程中,MMRazor 也针对蒸馏的基础模块做了一波升级,新增了 Recorder 和 Delivery 模块;通过 Recorder 可以获取任意位置的输入输出;通过 Delivery 可以实现 teacher 和 student 间任意位置数据的相互传递。

同时,进一步扩充了知识蒸馏算法的数量,覆盖了 data-free distillation 、zero-shot distillation 和 self distillation 等算法类型。

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第4张图片

*MMRazor 知识蒸馏模块*

NAS & Pruning 统一接口

在 OpenMMLab 1.0 时期,MMRazor 将 NAS 和 Pruning 抽象为了 Mutator 和 Pruner,经过半年的开发迭代,发现二者间有大量的可复用代码。

此次适配 OpenMMLab 2.0 的过程中,将 NAS 和 Pruning 统一抽象为了 Mutator,Mutabtor 会控制多个 Mutable (比如 NAS 中的 SearchBlock、Pruning 中的通道数),对于复杂的算法(结构和通道都会变化),通过多种 Mutator 的协作来完成复杂的形变操作。

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第5张图片

NAS & Pruning 调用关系

重点模块

Recorder

Recorder 是此次适配 OpenMMLab 2.0 的过程中新增的模块,用来自动地获取模型中任意位置的输入、输出数据。在使用的过程中可以同时使用多个 Recorder,一个 Recorder 只负责记录一个位置的数据,多个 Recorder 通过 Recorder Manager 统一管理。

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第6张图片

Recorder Pipeline

通过 Recorder,能够获得以下几种类型的中间结果,覆盖了蒸馏算法中的绝大多数场景:

  • Pytorch Module Inputs / Outputs
  • Pytorch Parameter
  • Python Function Inputs / Outputs
  • Python Class Method Inputs / Outputs

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第7张图片

Recorder 示例

Delivery

Delivery 也是此次适配 OpenMMLab 2.0 的过程中新增的模块,用来处理知识蒸馏算法中, student 和 teacher 间的数据传递(比如目标检测蒸馏中有时会将 teacher 的 label assign 结果赋给 student)。

和 Recorder 类似,在使用的过程中可以同时使用多个 Delivery,一个 Delivery 只负责记录一个位置的数据,多个 Delivery 通过 Delivery Manager 统一管理。

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第8张图片

Delivery Pipeline

结合开发过程的实际需求,目前 MMRazor 中支持以下两种 Delivery:

  • Python Function Outputs
  • Python Class Method Outputs

更多细节,可以移步 MMRazor delivery 文档 https://github.com/open-mmlab/mmrazor/blob/dev-1.x/docs/en/advanced_guides/delivery.md

Pruning Tracer

相较于 Pruning 算法本身,Pruning 算法的开发过程显得更加复杂,通道之间有很多依赖关系,一个算法想用在某一个模型上,往往需要重新开发一个模型,将各种通道的相关操作嵌入到 forward 的过程中。这个过程需要开发者对目标模型很熟悉,但即使开发者足够熟悉,开发的过程中也很容易出错,下图简单列举了一些常见的通道依赖:

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第9张图片

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第10张图片

为了能让剪枝算法能够在不同模型间快速切换,MMRazor 新增了能够自动解析通道依赖的 Pruning Tracer,只需 10 行以内代码,就可以快速获得通道依赖关系,配合 MMRazor 中的 Channel Mutator,可以自动化地根据通道依赖对模型进行改写,将剪枝算法的逻辑自动地嵌入到模型中。

更多细节,可以移步 MMRazor Pruning 文档 https://github.com/open-mmlab/mmrazor/blob/dev-1.x/docs/en/user_guides/pruning_user_guide.md

from mmrazor.models.mutators import ChannelMutator
from torchvision.models import mobilenet_v2

model = mobilenet_v2()
mutator = ChannelMutator(parse_cfg=dict(type='PruneTracer'))
mutator.prepare_from_supernet(model)

print(mutator.mutable_units[0])
# SequentialMutableChannelUnit(
#   name=features.0.0_(0, 32)_32
#   (output_related): ModuleList(
#     (0): Channel(features.0.0, index=(0, 32), is_output_channel=true, )
#     (1): Channel(features.0.1, index=(0, 32), is_output_channel=true, )
#     ...
#   )
#   (input_related): ModuleList(
#     (0): Channel(features.0.1, index=(0, 32), is_output_channel=false, )
#     (1): Channel(features.1.conv.0.0, index=(0, 32), is_output_channel=false, )
#     ...
#   )
#   (mutable_channel): SquentialMutableChannel(num_channels=32, activated_channels=32)
# )

总结

此次 OpenMMLab 2.0 的架构升级,打通了 MMRazor 与其他算法库,开发者可以使用 MMRazor 开发各种视觉任务的模型压缩算法,我们也在持续丰富算法种类,争取覆盖更多的算法库。

除了模型压缩算法外,MMRazor 还新增了一系列可以简化模型压缩算法开发的组件,用户也可以灵活地在自己的项目中(非 OpenMMLab 生态)使用这些组件,快速开发属于自己的模型压缩算法。

维护计划

因为 MMRazor 1.0 相较于 0.x 版本变动较大,我们也采取了较为保守的维护策略,确保熟悉了 0.x 版本的用户有充足的时间迁移至新版本。维护计划如下图所示,整体按时间会分为公测期、兼容期和维护期。整体按时间会分为公测期、兼容期和维护期。

公测期(2022/9/1 - 2022/12/31)

  • 公测版代码发布在 1.x 分支,开发分支为 dev-1.x新特性和新模型主要在 1.x 上发布
  • 默认主分支 master 为 MMRazor 0.x 版本,保持对代码的维护和修复,响应用户需求,不会增加主要新功能

兼容期(2023/1/1 - 2023/12/31)

  • 默认主分支切换为 1.x 分支,开发分支仍为 dev-1.x
  • 进一步完善从 MMRazor 0.x 到 MMRazor 1.0 的迁移指引,协助用户进行统一迁移
  • 0.x 系列版本主分支切换到 0.x,保持对旧版本的维护,并确保不引进破坏旧版本兼容性的改动

维护期(2024 起)

  • 旧版本进入维护期,不再进行新功能支持,主要投入到新版本的开发迭代中

模型压缩算法库 MMRazor 全面升级,更灵活、更自动_第11张图片

你可能感兴趣的:(新闻速递,人工智能)