作者:机器学习与边缘侧AI首席工程师Sheila Seidel
2025年10月30日
为何缩减模型大小至关重要
缩小AI模型并非锦上添花的优化,而是为边缘设备赋予强大实时智能的必要之举。无论是智能手机、可穿戴设备还是嵌入式系统,都受限于内存、算力和电池容量方面的约束。在设备端提供AI体验,意味着要在不牺牲模型精度与性能的前提下,实现超低延迟和超低功耗 。
而这正是模型压缩技术的用武之地。非结构化剪枝是常用的初始手段(移除单个低重要性权重,对整体架构影响极小),但这种方法往往会产生不规则的稀疏模式,且因掩码开销导致运行时效率低下。因此,ADI团队将重点转向结构化剪枝,真正降低内存占用和计算成本。更重要的是,经由这种方法得到的模型不仅体积更小,还能更轻松地部署在资源受限的边缘设备上。
结构化剪枝:提升模型效率的更优路径
结构化剪枝为在资源受限平台上部署更小模型提供了切实可行的方案。从基础形式来看,这项技术会移除完整的单元,例如通道、滤波器,甚至整个层。更重要的是,与非结构化剪枝需要追踪非零权重的位置不同,结构化剪枝能生成更简洁、更小的权重矩阵,优化与部署皆事半功倍。
然而,剪枝并非简单地删除模型的部分组件。移除某一层的参数可能会牵一发而动全身,尤其是在循环神经网络(RNN)这类紧密耦合的架构中。处理好各种依赖关系,是确保剪枝后的模型保持功能性和有效性的关键所在。
超越内置工具
尽管PyTorch通过torch.nn.utils.prune提供了基础剪枝工具,但这些工具存在局限:仅能应用掩码,无法缩减模型的实际大小,也无法处理层间的依赖关系。
为解决这一问题,我们一直在使用 torch-pruning开源库,它能从模型中物理移除冗余组件。与PyTorch的内置工具不同,torch-pruning会构建网络的依赖关系图,以确定各层之间的交互方式。这使得结构化剪枝能够遵循模型的架构逻辑进行。
实际应用:语音降噪和语音活动检测
我们近期承接了一个语音降噪项目,借助torch-pruning工具对卷积-GRU模型进行压缩。Conv2D层的剪枝过程直截了当,在不影响性能的前提下, 将乘加运算(MAC)量减少了25%至30% 。
但是,GRU层的剪枝却遇阻受挫。PyTorch的GRU由C++实现,与torch-pruning的依赖关系图兼容性不甚理想。此外,由于GRU会将其隐藏状态反馈回自身,在默认情况下,无法对其输入进行剪枝。
为此,我们的解决方案是基于PyTorch标准层(如nn.Linear)和元素级运算,从零开始构建可剪枝的GRU。我们还引入了恒等适配层,将隐藏状态的大小与GRU输入解耦,从而实现安全剪枝。
具体过程如下:
- 用自定义的可剪枝GRU替换原生GRU,并复制现有权重。
- 借助torch-pruning执行结构化剪枝。
- 将可剪枝GRU转换回原生GRU。
- 对模型进行微调以恢复准确度。
- 部署经过缩减后的轻量模型。
这种方法除了通过卷积层剪枝来减少乘加运算量外,更使模型参数减少了30%。
在我们的语音活动检测(VAD)项目中,模型架构以卷积层为核心。我们将原本含8万参数的模型剪枝至2千参数以内 ,同时保持了优异性能。上述成果的关键在于:对后续层进行精细剪枝,同时保留早期层的丰富特性。
开源贡献和未来展望
我们的自定义GRU剪枝解决方案现已作为开源拉取请求提交至torch-pruning库: PR #515。这项工作为其他循环单元(如LSTM和精简版RNN变体)的剪枝铺平了道路。
展望未来,我们正探索智能剪枝工作流,力求让结构化剪枝更易于使用,自适应更多样的模型架构。
结构化剪枝并非仅停留于研究层面的好奇探索,而是一种实用工具,能让强大的AI模型更高效、更负责任地运行在日常设备上。