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

news2025/7/24 9:05:29

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 支持算法类型

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

新版本特点

全面支持 OpenMMLab

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

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

在这里插入图片描述

Config 更加简洁

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

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

在这里插入图片描述

知识蒸馏全新升级

相较于知识蒸馏算法本身，知识蒸馏的算法实现会更加繁琐一些。很多知识蒸馏算法都是在围绕如何设计知识蒸馏的损失函数，在实现过程中，为了获取到要计算蒸馏损失的特征图，往往要通过改写模型 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 知识蒸馏模块*

NAS & Pruning 统一接口

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

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

在这里插入图片描述

NAS & Pruning 调用关系

重点模块

Recorder

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

在这里插入图片描述

Recorder Pipeline

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

Pytorch Module Inputs / Outputs

Pytorch Parameter

Python Function Inputs / Outputs

Python Class Method Inputs / Outputs

在这里插入图片描述

Recorder 示例

Delivery

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

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

在这里插入图片描述

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 新增了能够自动解析通道依赖的 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)
# )