[论文解读] GNN-RL Compression: Topology-Aware Network Pruning using Multi-stage Graph Embedding and Reinforcement Learning.

Sixing Yu, Arya Mazaheri+1 more|arXiv (Cornell University)|Feb 5, 2021

Advanced Graph Neural Networks参考文献 36被引用 8

一句话总结

该论文提出 GNN-RL Compression，一种拓扑感知的通道剪枝方法，利用多阶段图神经网络（GNNs）从深度神经网络（DNNs）中提取结构信息，并应用强化学习（RL）在 FLOPs 约束下优化压缩策略。该方法在 VGG-16 上实现了高达 80% 的 FLOPs 减少，在 ShuffleNet-v1 上相比最先进方法准确率最高提升 1.84%，且在剪枝后的 VGG-16 上实现了 1.38 倍的推理速度提升和 141 MB 的 GPU 显存减少。

ABSTRACT

Model compression is an essential technique for deploying deep neural networks (DNNs) on power and memory-constrained resources. However, existing model-compression methods often rely on human expertise and focus on parameters' local importance, ignoring the rich topology information within DNNs. In this paper, we propose a novel multi-stage graph embedding technique based on graph neural networks (GNNs) to identify the DNNs' topology and use reinforcement learning (RL) to find a suitable compression policy. We performed resource-constrained (i.e., FLOPs) channel pruning and compared our approach with state-of-the-art compression methods using over-parameterized DNNs (e.g., ResNet and VGG-16) and mobile-friendly DNNs (e.g., MobileNet and ShuffleNet). We evaluated our method on various models from typical to mobile-friendly networks, such as ResNet family, VGG-16, MobileNet-v1/v2, and ShuffleNet. The results demonstrate that our method can prune dense networks (e.g., VGG-16) by up to 80% of their original FLOPs. More importantly, our method outperformed state-of-the-art methods and achieved a higher accuracy by up to 1.84% for ShuffleNet-v1. Furthermore, following our approach, the pruned VGG-16 achieved a noticeable 1.38$ imes$ speed up and 141 MB GPU memory reduction.

研究动机与目标

• 为解决现有模型压缩方法忽略 DNN 中结构拓扑且依赖局部参数重要性的局限性。
• 开发一种拓扑感知的剪枝框架，以捕捉深度神经网络中的全局结构依赖关系。
• 利用强化学习在严格的 FLOPs 约束下优化压缩策略。
• 在从过参数化的模型（如 ResNet 和 VGG-16）到轻量级网络（如 MobileNet 和 ShuffleNet）等多种架构上评估该方法。
• 在资源受限环境下，实现相比最先进压缩技术更高的准确率和效率。

提出的方法

• 采用基于图神经网络（GNNs）的多阶段图嵌入技术，以建模 DNN 中各层之间的分层和结构关系。
• 基于 GNN 的嵌入能够捕捉网络架构的拓扑感知表示，从而在参数量之外更准确地理解结构重要性。
• 训练一个强化学习智能体，根据 GNN 学习到的拓扑嵌入选择通道剪枝策略，以优化准确率和 FLOPs 约束。
• RL 智能体在序列决策过程中运行，其中动作对应于在每一层中剪枝特定通道，以满足目标 FLOPs 预算。
• 该方法端到端应用于多种 DNN，包括 ResNet、VGG-16、MobileNet-v1/v2 和 ShuffleNet，剪枝由结构信息和性能反馈共同指导。
• 训练流程将基于 GNN 的拓扑编码与基于 RL 的策略优化相结合，实现自动化的、拓扑感知的压缩。

实验结果

研究问题

• RQ1 基于 GNN 的方法能否有效捕捉深度神经网络的结构拓扑，以提升剪枝决策？
• RQ2 与基于参数重要性的传统剪枝方法相比，使用 GNN 和 RL 的拓扑感知剪枝在准确率和效率方面表现如何？
• RQ3 所提方法在多种 DNN 架构上能在多大程度上减少 FLOPs，同时保持或提升模型准确率？
• RQ4 在实际部署场景中，该方法能否在移动友好型模型上实现显著的速度提升和显存减少？
• RQ5 GNN-RL 压缩框架在不同类型的网络（从过参数化到轻量级架构）上具有多强的鲁棒性？

主要发现

• 所提方法在 VGG-16 上实现了高达 80% 的 FLOPs 减少，同时保持高准确率，展现出在密集网络上的强大压缩能力。
• 在 ShuffleNet-v1 上，该方法相比最先进方法准确率最高提升 1.84%，凸显其在移动场景下的有效性。
• 剪枝后的 VGG-16 模型推理时间实现 1.38 倍加速，GPU 显存使用减少 141 MB，表明其具备显著的实际效率优势。
• 该方法在多种架构（包括 ResNet、VGG-16、MobileNet-v1/v2 和 ShuffleNet）上泛化良好，表现出一致的性能提升。
• 通过利用 GNN 的拓扑感知表示和基于 RL 的策略学习，该方法在模型大小、速度和准确率之间的权衡上优于现有技术。