[论文解读] Compressing Vision Transformers for Low-Resource Visual Learning

Eric Youn, Sai Mitheran J.+2 more|arXiv (Cornell University)|Sep 5, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

本文提出了一种用于视觉变换器的混合压缩框架，以实现在低资源边缘设备（如NVIDIA Jetson Nano，4GB内存）上的高效部署，结合结构化剪枝、知识蒸馏和训练后量化。该方法在极低延迟和内存占用下实现了与Swin Transformer相当的精度，使在无人机灾害监测等资源受限场景下实现实时语义分割成为可能。

ABSTRACT

Vision transformer (ViT) and its variants have swept through visual learning leaderboards and offer state-of-the-art accuracy in tasks such as image classification, object detection, and semantic segmentation by attending to different parts of the visual input and capturing long-range spatial dependencies. However, these models are large and computation-heavy. For instance, the recently proposed ViT-B model has 86M parameters making it impractical for deployment on resource-constrained devices. As a result, their deployment on mobile and edge scenarios is limited. In our work, we aim to take a step toward bringing vision transformers to the edge by utilizing popular model compression techniques such as distillation, pruning, and quantization. Our chosen application environment is an unmanned aerial vehicle (UAV) that is battery-powered and memory-constrained, carrying a single-board computer on the scale of an NVIDIA Jetson Nano with 4GB of RAM. On the other hand, the UAV requires high accuracy close to that of state-of-the-art ViTs to ensure safe object avoidance in autonomous navigation, or correct localization of humans in search-and-rescue. Inference latency should also be minimized given the application requirements. Hence, our target is to enable rapid inference of a vision transformer on an NVIDIA Jetson Nano (4GB) with minimal accuracy loss. This allows us to deploy ViTs on resource-constrained devices, opening up new possibilities in surveillance, environmental monitoring, etc. Our implementation is made available at https://github.com/chensy7/efficient-vit.

研究动机与目标

• 在内存受限、电池供电的边缘设备（如配备NVIDIA Jetson Nano，4GB内存的无人机）上实现视觉变换器的部署。
• 在最小化模型尺寸和推理延迟的同时，保持与Swin Transformer相当的分割精度。
• 开发一个实用的端到端压缩流水线，结合结构化剪枝、蒸馏和训练后量化，以实现边缘部署。
• 在严格内存和延迟约束下，基于LPCV数据集对框架在真实灾害场景解析中的表现进行评估。

提出的方法

• 对注意力头和线性层中的神经元应用结构化剪枝，减少参数量，同时不改变计算图结构。
• 采用多阶段蒸馏策略：在预测头进行logits知识迁移，在浅层进行特征模仿，在深层采用基于生成的蒸馏。
• 使用训练后量化（PTQ）以减小模型大小和内存占用，尽管PyTorch对量化CUDA操作的支持有限。
• 在剪枝后结合蒸馏与迭代微调，以稳定性能并提升泛化能力。
• 使用更强的Swin-v2-T模型作为教师，指导小型MobileViT学生网络的知识迁移。
• 以平均交并比（mIoU）为主要指标，在LPCV数据集上评估压缩后的模型。

Figure 1: Our framework for edge-oriented, rapid on-device disaster scene parsing

实验结果

研究问题

• RQ1 在边缘设备上，对注意力头和线性层权重进行结构化剪枝是否能有效减小模型尺寸，同时避免显著的精度下降？
• RQ2 采用多阶段监督（logits、特征、生成）的知识蒸馏在将大型教师模型的性能迁移到紧凑的学生ViT模型方面有多有效？
• RQ3 在Jetson Nano上，训练后量化在多大程度上能减少内存使用和延迟，同时保持分割精度？
• RQ4 在真实灾害场景中，类别不平衡和细粒度视觉差异（如泥流与土路）如何影响压缩模型的性能？
• RQ5 与一次性剪枝相比，采用迭代剪枝结合微调是否能显著提升压缩效率？

主要发现

• 最终压缩模型在仅4GB系统内存的条件下，于LPCV数据集上达到了与Swin Transformer相当的精度，满足边缘部署约束。
• 将注意力头从4个减少到2个对性能影响极小，显著减少了参数量，同时保持了高mIoU。
• 训练后量化显著降低了内存使用，但因PyTorch对量化CUDA操作的支持有限，进一步优化受到限制。
• 模型在主要类别上表现良好，但在泥流、土路等细粒度、低频类别上表现不佳，凸显了类别不平衡的挑战。
• 与一次性剪枝相比，迭代剪枝仅带来微小改进，表明对该特定架构和压缩目标而言，收益有限。
• 该框架成功实现在Jetson Nano上的实时语义分割，证明了其在无人机灾害监测中的可行性。

Figure 2: Samples from the LPCV dataset (CVPR 2023 challenge)