使用轻量级背景感知视觉Transformer的令牌剪枝

高运行时内存和高延迟对Vision Transformer的训练和推理构成了显著限制，尤其是在边缘设备上。令牌剪枝根据每个令牌的重要性标准减少输入到ViT的令牌数量。我们提出了一种背景感知视觉Transformer（BAViT）模型，作为DETR/YOLOS等目标检测模型的预处理模块，旨在通过一种新颖的方法识别图像中的背景令牌来减少运行时内存并提高吞吐量。在将图像输入基于ViT的目标检测器之前，可以完全或部分剪枝背景令牌。我们利用分割图和/或边界框标注提供的语义信息来训练几层ViT，以将令牌分类为前景或背景。使用2层和10层的BAViT，在VOC数据集上，背景和前景令牌的分类准确率分别达到75%和88%；在COCO数据集上，分别达到71%和80%。我们展示了将2层BAViT-small模型作为YOLOS的预处理器，可以在没有稀疏微调的情况下，将吞吐量提高30%-40%，同时mAP下降3%；在有稀疏微调的情况下，mAP下降2%。我们的方法专门针对边缘AI应用场景。