FastViT: 快速混合视觉Transformer的结构重参数化

FastViT简介

FastViT(Fast Hybrid Vision Transformer)是由苹果公司研究人员提出的一种新型视觉Transformer模型。该模型通过结构重参数化技术,在保持高准确率的同时大幅提升了推理速度,特别适合在移动设备等计算资源受限的场景下使用。

FastViT的核心思想是结合了CNN和Transformer的优点,采用混合架构设计。在模型的早期阶段使用轻量级CNN提取局部特征,后期阶段则使用Transformer进行全局建模。这种设计既保留了CNN的效率,又利用了Transformer的强大表达能力。

如上图所示,FastViT系列模型在ImageNet分类任务上取得了优秀的准确率和速度权衡。与其他流行的轻量级CNN和Transformer模型相比,FastViT在相同延迟下可以获得更高的Top-1准确率。

模型架构

FastViT的整体架构包含以下几个关键组件:

1. 轻量级CNN主干网络:采用MobileNetV2的倒置残差结构,用于提取局部特征。
2. 轻量级Transformer块:使用自注意力机制进行全局特征建模。
3. 结构重参数化:通过重参数化技术,将训练时的复杂结构转换为推理时的简单结构,提高推理效率。
4. 渐进式通道扩展:随着网络深度增加,逐步增加特征通道数,增强模型表达能力。
5. 跳跃连接:采用残差连接,有利于梯度传播和特征重用。

这种混合架构设计使FastViT能够在保持高准确率的同时,显著降低计算复杂度和内存占用。

结构重参数化

结构重参数化是FastViT实现高效推理的关键技术。其核心思想是:在训练阶段使用复杂的多分支结构,而在推理阶段将多分支结构等效转换为单一的线性层。

具体来说,FastViT在训练时采用了类似MobileNetV2的倒置残差结构,包含1x1卷积、3x3深度可分离卷积和跳跃连接。在推理阶段,这些分支可以通过数学变换合并为一个等效的3x3卷积,大大简化了计算图,提高了推理效率。

这种技术使FastViT能够在训练时保持良好的表达能力和优化特性,同时在推理时获得接近轻量级CNN的高效率。

模型性能

FastViT在ImageNet-1K数据集上进行了广泛的实验评估。下面列出了几个代表性变体的性能数据:

所有模型都在iPhone 12 Pro上使用CoreML进行了延迟测试。可以看到,FastViT系列模型在1-5ms的延迟范围内实现了76%-84%的Top-1准确率,展现了优秀的速度-准确率权衡。

特别值得一提的是,FastViT-SA12模型在1.6ms延迟下达到了80.9%的准确率,这一性能在当前移动端视觉模型中处于领先水平。

知识蒸馏

为进一步提升模型性能,研究人员还尝试了知识蒸馏技术。具体做法是使用大型模型(如EfficientNetV2-L)作为教师模型,指导FastViT模型的训练。

实验结果表明,知识蒸馏可以显著提升FastViT的性能。例如:

• FastViT-T8: 准确率从76.2%提升到77.2%
• FastViT-SA36: 准确率从83.6%提升到84.2%

这说明FastViT模型具有进一步提升的潜力,通过蒸馏等技术可以进一步缩小与大型模型的性能差距。

使用指南

FastViT模型的使用非常简便。研究人员提供了预训练权重和完整的训练、评估代码。下面是一个简单的使用示例:

import torch
from timm.models import create_model
from models.modules.mobileone import reparameterize_model

# 加载预训练模型
model = create_model("fastvit_t8")
checkpoint = torch.load('/path/to/checkpoint.pth.tar')
model.load_state_dict(checkpoint['state_dict'])

# 转换为推理模式
model.eval()      
model_inf = reparameterize_model(model)

# 使用模型进行推理
...

此外,研究人员还提供了将PyTorch模型转换为CoreML格式的脚本,方便在iOS设备上部署。

总结与展望

FastViT作为一种新型的混合视觉Transformer模型,在移动端视觉任务中展现出了巨大的潜力。其主要优势包括:

1. 优秀的速度-准确率权衡,特别适合移动设备等计算资源受限的场景。

2. 灵活的模型系列,覆盖了76%-84%的准确率范围,可根据实际需求选择合适的变体。

3. 简单易用的API和完善的工具链,方便研究人员和开发者使用和部署。

4. 开源的代码和预训练模型,有利于社区进一步改进和应用。

未来,FastViT还有进一步发展的空间。可能的研究方向包括:

• 探索更高效的混合架构设计
• 改进结构重参数化技术,进一步提升推理效率
• 将FastViT应用到更多的计算机视觉任务中,如目标检测、语义分割等
• 针对特定硬件平台进行优化,发挥极致性能

总的来说,FastViT为移动端视觉AI的发展提供了新的思路和工具,相信会推动更多高效视觉AI应用的落地。研究人员和开发者可以基于FastViT开发各种创新的移动端视觉应用,为用户带来更智能、更便捷的AI体验。