从头训练 N:M 细粒度结构化稀疏神经网络
计算机视觉与模式识别
2021-04-20 v2 硬件体系结构
摘要
深度神经网络(DNNs)中的稀疏性已被广泛研究,以在资源受限环境下压缩和加速模型。它通常可分为非结构化细粒度稀疏(将神经网络中分布的多个独立权重置零)和结构化粗粒度稀疏(剪枝神经网络的子网络块)。细粒度稀疏可实现高压缩比,但对硬件不友好,因此获得的加速有限。另一方面,粗粒度稀疏无法同时在现代 GPU 上实现明显的加速和良好的性能。在本文中,我们首次研究从头训练 N:M 细粒度结构化稀疏网络,该网络在专门设计的 GPU 上可同时保持非结构化细粒度稀疏和结构化粗粒度稀疏的优势。具体而言,2:4 稀疏网络在 Nvidia A100 GPU 上可实现 2 倍加速且无性能下降。此外,我们提出了一种新颖而有效的成分——稀疏精炼直通估计器(SR-STE),以缓解优化过程中由朴素 STE 计算的近似梯度带来的负面影响。我们还定义了一个度量指标——稀疏架构散度(SAD),用于衡量训练过程中稀疏网络拓扑的变化。最后,我们用 SAD 证明了 SR-STE 的优势,并通过在各种任务上的综合实验展示了 SR-STE 的有效性。源代码和模型可在 https://github.com/NM-sparsity/NM-sparsity 获取。
引用
@article{arxiv.2102.04010,
title = {Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch},
author = {Aojun Zhou and Yukun Ma and Junnan Zhu and Jianbo Liu and Zhijie Zhang and Kun Yuan and Wenxiu Sun and Hongsheng Li},
journal= {arXiv preprint arXiv:2102.04010},
year = {2021}
}
备注
ICLR2021