English

Topology-Aware Network Pruning using Multi-stage Graph Embedding and Reinforcement Learning

Computer Vision and Pattern Recognition 2022-07-04 v2 Machine Learning

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 DNN topologies 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 model compression methods. We evaluated our method on various models from typical to mobile-friendly networks, such as ResNet family, VGG-16, MobileNet-v1/v2, and ShuffleNet. Results show that our method can achieve higher compression ratios with a minimal fine-tuning cost yet yields outstanding and competitive performance.

Keywords

Cite

@article{arxiv.2102.03214,
  title  = {Topology-Aware Network Pruning using Multi-stage Graph Embedding and Reinforcement Learning},
  author = {Sixing Yu and Arya Mazaheri and Ali Jannesari},
  journal= {arXiv preprint arXiv:2102.03214},
  year   = {2022}
}

Comments

Accepted at ICML 2022 Long presentation

R2 v1 2026-06-23T22:52:33.600Z