FastViT简介
FastViT(Fast Hybrid Vision Transformer)是由苹果公司研究人员提出的一种新型视觉Transformer模型。该模型通过结构重参数化技术,在保持高准确率的同时大幅提升了推理速度,特别适合在移动设备等计算资源受限的场景下使用。
FastViT的核心思想是结合了CNN和Transformer的优点,采用混合架构设计。在模型的早期阶段使用轻量级CNN提取局部特征,后期阶段则使用Transformer进行全局建模。这种设计既保留了CNN的效率,又利用了Transformer的强大表达能力。
如上图所示,FastViT系列模型在ImageNet分类任务上取得了优秀的准确率和速度权衡。与其他流行的轻量级CNN和Transformer模型相比,FastViT在相同延迟下可以获得更高的Top-1准确率。
模型架构
FastViT的整体架构包含以下几个关键组件:
- 轻量级CNN主干网络:采用MobileNetV2的倒置残差结构,用于提取局部特征。
- 轻量级Transformer块:使用自注意力机制进行全局特征建模。
- 结构重参数化:通过重参数化技术,将训练时的复杂结构转换为推理时的简单结构,提高推理效率。
- 渐进式通道扩展:随着网络深度增加,逐步增加特征通道数,增强模型表达能力。
- 跳跃连接:采用残差连接,有利于梯度传播和特征重用。
这种混合架构设计使FastViT能够在保持高准确率的同时,显著降低计算复杂度和内存占用。
结构重参数化
结构重参数化是FastViT实现高效推理的关键技术。其核心思想是:在训练阶段使用复杂的多分支结构,而在推理阶段将多分支结构等效转换为单一的线性层。
具体来说,FastViT在训练时采用了类似MobileNetV2的倒置残差结构,包含1x1卷积、3x3深度可分离卷积和跳跃连接。在推理阶段,这些分支可以通过数学变换合并为一个等效的3x3卷积,大大简化了计算图,提高了推理效率。
这种技术使FastViT能够在训练时保持良好的表达能力和优化特性,同时在推理时获得接近轻量级CNN的高效率。
模型性能
FastViT在ImageNet-1K数据集上进行了广泛的实验评估。下面列出了几个代表性变体的性能数据:
| 模型 | Top-1准确率 | 延迟(ms) |
|---|---|---|
| FastViT-T8 | 76.2% | 0.8 |
| FastViT-T12 | 79.3% | 1.2 |
| FastViT-S12 | 79.9% | 1.4 |
| FastViT-SA12 | 80.9% | 1.6 |
| FastViT-SA24 | 82.7% | 2.6 |
| FastViT-SA36 | 83.6% | 3.5 |
| FastViT-MA36 | 83.9% | 4.6 |
所有模型都在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模型,在移动端视觉任务中展现出了巨大的潜力。其主要优势包括:
-
优秀的速度-准确率权衡,特别适合移动设备等计算资源受限的场景。
-
灵活的模型系列,覆盖了76%-84%的准确率范围,可根据实际需求选择合适的变体。
-
简单易用的API和完善的工具链,方便研究人员和开发者使用和部署。
-
开源的代码和预训练模型,有利于社区进一步改进和应用。
未来,FastViT还有进一步发展的空间。可能的研究方向包括:
- 探索更高效的混合架构设计
- 改进结构重参数化技术,进一步提升推理效率
- 将FastViT应用到更多的计算机视觉任务中,如目标检测、语义分割等
- 针对特定硬件平台进行优化,发挥极致性能
总的来说,FastViT为移动端视觉AI的发展提供了新的思路和工具,相信会推动更多高效视觉AI应用的落地。研究人员和开发者可以基于FastViT开发各种创新的移动端视觉应用,为用户带来更智能、更便捷的AI体验。
