FasterViT：具有层次注意力的快速视觉 Transformer

FasterViT：具有层次注意力的快速视觉 Transformer 的官方 PyTorch 实现。

Ali Hatamizadeh、 Greg Heinrich、 Hongxu (Danny) Yin、 Andrew Tao、 Jose M. Alvarez、 Jan Kautz、 Pavlo Molchanov。

如需商业咨询，请访问我们的网站并提交表单：NVIDIA 研究许可

FasterViT 在不使用额外训练数据的情况下，在 Top-1 准确率和吞吐量方面实现了新的 SOTA 帕累托前沿！

我们引入了一种新的自注意力机制，称为层次注意力（HAT），通过学习跨窗口载体令牌来捕获短程和长程信息。

预览图

注意：请使用最新的 NVIDIA TensorRT 版本以享受优化的 FasterViT 操作带来的好处。

💥 新闻 💥

[2024.04.02] 🔥 更新后的论文现已在 arXiv 上发布！
[2024.01.24] 🔥🔥🔥 使用 MOTRv2 + FasterViT 的目标跟踪现已开源（链接）！
[2024.01.17] 🔥🔥🔥 FasterViT 论文已被 ICLR 2024 接收！
[2023.10.14] 🔥🔥 我们已添加了使用 DINO 的 FasterViT 目标检测仓库！
[2023.08.24] 🔥 FasterViT Keras 模型及预训练权重已在 keras_cv_attention_models 中发布！
[2023.08.20] 🔥🔥 我们已添加了各种分辨率的 ImageNet-21K SOTA 预训练模型！
[2023.07.20] 我们已创建官方 NVIDIA FasterViT HuggingFace 页面。
[2023.07.06] FasterViT 检查点现在也可在 HuggingFace 上访问！
[2023.07.04] 现在可以使用一行代码导入 ImageNet 预训练的 FasterViT 模型。请安装最新的 FasterViT pip 包以使用此功能（也支持任意分辨率的 FasterViT 模型）。
[2023.06.30] 我们进一步提高了 FasterViT 模型在 TensorRT 上的吞吐量，不同模型平均提升了 10-15%。请使用最新的 NVIDIA TensorRT 版本以获得这些吞吐量性能提升。
[2023.06.29] 任意分辨率的 FasterViT 模型现在可以从预训练的 ImageNet 分辨率（224 x 224）模型初始化。
[2023.06.18] 我们已发布 FasterViT pip 包！
[2023.06.17] 任意分辨率 FasterViT 模型现已可用！该模型可用于各种应用，如检测和分割或任意输入图像分辨率的高分辨率微调。
[2023.06.09] 🔥🔥 我们已发布源代码和 ImageNet-1K FasterViT 模型！

快速开始

目标检测

请查看 FasterViT 目标检测仓库，其中使用了 DINO：具有改进去噪锚框的端到端目标检测 DETR，以获取更多详细信息。

分类

我们可以使用一行代码导入预训练的 FasterViT 模型。首先，可以简单地安装 FasterViT：

pip install fastervit

注意：如果您已经安装了该包，请将其升级到 fastervit>=0.9.8 以使用预训练权重。

可以按如下方式创建具有默认超参数的预训练 FasterViT 模型：

>>> from fastervit import create_model

# 定义分辨率为 224 x 224 的 fastervit-0 模型

>>> model = create_model('faster_vit_0_224', 
                          pretrained=True,
                          model_path="/tmp/faster_vit_0.pth.tar")

model_path 用于设置下载模型的目录。

我们还可以通过传入一个虚拟输入图像来简单测试模型。输出是 logits：

>>> import torch

>>> image = torch.rand(1, 3, 224, 224)
>>> output = model(image) # torch.Size([1, 1000])

我们还可以使用任意分辨率的 FasterViT 模型来适应任意图像分辨率。在下面的例子中，我们定义了一个输入分辨率为 576 x 960、第 3 和第 4 阶段的窗口大小分别为 12 和 6、载体令牌大小为 2、嵌入维度为 64 的任意分辨率 FasterViT-0 模型：

>>> from fastervit import create_model

# 定义分辨率为 576 x 960 的任意分辨率 FasterViT-0 模型
>>> model = create_model('faster_vit_0_any_res', 
                          resolution=[576, 960],
                          window_size=[7, 7, 12, 6],
                          ct_size=2,
                          dim=64,
                          pretrained=True)

请注意，上述模型是从原始分辨率为 224 x 224 的 ImageNet 预训练 FasterViT 初始化的。因此，由于我们添加了新层（例如添加新的载体令牌等），可能会出现缺失键和不匹配的情况。

我们可以通过传入一个虚拟输入图像来测试模型。输出是 logits：

>>> import torch

>>> image = torch.rand(1, 3, 576, 960)
>>> output = model(image) # torch.Size([1, 1000])

结果 + 预训练模型

ImageNet-1K

FasterViT ImageNet-1K 预训练模型

<table> <tr> <th>名称</th> <th>Acc@1(%)</th> <th>Acc@5(%)</th> <th>吞吐量(图像/秒)</th> <th>分辨率</th> <th>参数量(M)</th> <th>FLOPs(G)</th> <th>下载</th> </tr> <tr> <td>FasterViT-0</td> <td>82.1</td> <td>95.9</td> <td>5802</td> <td>224x224</td> <td>31.4</td> <td>3.3</td> <td><a href="https://drive.google.com/uc?export=download&id=1twI2LFJs391Yrj8MR4Ui9PfrvWqjE1iB">模型</a></td> </tr> <tr> <td>FasterViT-1</td> <td>83.2</td> <td>96.5</td> <td>4188</td> <td>224x224</td> <td>53.4</td> <td>5.3</td> <td><a href="https://drive.google.com/uc?export=download&id=1r7W10n5-bFtM3sz4bmaLrowN2gYPkLGT">模型</a></td> </tr> <tr> <td>FasterViT-2</td> <td>84.2</td> <td>96.8</td> <td>3161</td> <td>224x224</td> <td>75.9</td> <td>8.7</td> <td><a href="https://drive.google.com/uc?export=download&id=1n_a6s0pgi0jVZOGmDei2vXHU5E6RH5wU">模型</a></td> </tr> <tr> <td>FasterViT-3</td> <td>84.9</td> <td>97.2</td> <td>1780</td> <td>224x224</td> <td>159.5</td> <td>18.2</td> <td><a href="https://drive.google.com/uc?export=download&id=1tvWElZ91Sia2SsXYXFMNYQwfipCxtI7X">模型</a></td> </tr> <tr> <td>FasterViT-4</td> <td>85.4</td> <td>97.3</td> <td>849</td> <td>224x224</td> <td>424.6</td> <td>36.6</td> <td><a href="https://drive.google.com/uc?export=download&id=1gYhXA32Q-_9C5DXel17avV_ZLoaHwdgz">模型</a></td> </tr> <tr> <td>FasterViT-5</td> <td>85.6</td> <td>97.4</td> <td>449</td> <td>224x224</td> <td>975.5</td> <td>113.0</td> <td><a href="https://drive.google.com/uc?export=download&id=1mqpai7XiHLr_n1tjxjzT8q369xTCq_z-">模型</a></td> </tr> <tr> <td>FasterViT-6</td> <td>85.8</td> <td>97.4</td> <td>352</td> <td>224x224</td> <td>1360.0</td> <td>142.0</td> <td><a href="https://drive.google.com/uc?export=download&id=12jtavR2QxmMzcKwPzWe7kw-oy34IYi59">模型</a></td> </tr> </table>

ImageNet-21K

FasterViT ImageNet-21K预训练模型（在ImageNet-1K上微调）

<table> <tr> <th>名称</th> <th>Acc@1(%)</th> <th>Acc@5(%)</th> <th>分辨率</th> <th>参数量(M)</th> <th>FLOPs(G)</th> <th>下载</th> </tr> <tr> <td>FasterViT-4-21K-224</td> <td>86.6</td> <td>97.8</td> <td>224x224</td> <td>271.9</td> <td>40.8</td> <td><a href="https://huggingface.co/ahatamiz/FasterViT/resolve/main/fastervit_4_21k_224_w14.pth.tar">模型</a></td> </tr> <tr> <td>FasterViT-4-21K-384</td> <td>87.6</td> <td>98.3</td> <td>384x384</td> <td>271.9</td> <td>120.1</td> <td><a href="https://huggingface.co/ahatamiz/FasterViT/resolve/main/fastervit_4_21k_384_w24.pth.tar">模型</a></td> </tr> <tr> <td>FasterViT-4-21K-512</td> <td>87.8</td> <td>98.4</td> <td>512x512</td> <td>271.9</td> <td>213.5</td> <td><a href="https://huggingface.co/ahatamiz/FasterViT/resolve/main/fastervit_4_21k_512_w32.pth.tar">模型</a></td> </tr> <tr> <td>FasterViT-4-21K-768</td> <td>87.9</td> <td>98.5</td> <td>768x768</td> <td>271.9</td> <td>480.4</td> <td><a href="https://huggingface.co/ahatamiz/FasterViT/resolve/main/fastervit_4_21k_768_w48.pth.tar">模型</a></td> </tr> </table>

FasterViT-4在ImageNet-21K上的原始预训练模型权重也可以通过此链接下载。

鲁棒性（ImageNet-A - ImageNet-R - ImageNet-V2）

所有模型使用crop_pct=0.875。结果是通过在ImageNet-1K预训练模型上直接进行推理获得的，没有进行微调。

<table> <tr> <th>名称</th> <th>A-Acc@1(%)</th> <th>A-Acc@5(%)</th> <th>R-Acc@1(%)</th> <th>R-Acc@5(%)</th> <th>V2-Acc@1(%)</th> <th>V2-Acc@5(%)</th> </tr> <tr> <td>FasterViT-0</td> <td>23.9</td> <td>57.6</td> <td>45.9</td> <td>60.4</td> <td>70.9</td> <td>90.0</td> </tr> <tr> <td>FasterViT-1</td> <td>31.2</td> <td>63.3</td> <td>47.5</td> <td>61.9</td> <td>72.6</td> <td>91.0</td> </tr> <tr> <td>FasterViT-2</td> <td>38.2</td> <td>68.9</td> <td>49.6</td> <td>63.4</td> <td>73.7</td> <td>91.6</td> </tr> <tr> <td>FasterViT-3</td> <td>44.2</td> <td>73.0</td> <td>51.9</td> <td>65.6</td> <td>75.0</td> <td>92.2</td> </tr> <tr> <td>FasterViT-4</td> <td>49.0</td> <td>75.4</td> <td>56.0</td> <td>69.6</td> <td>75.7</td> <td>92.7</td> </tr> <tr> <td>FasterViT-5</td> <td>52.7</td> <td>77.6</td> <td>56.9</td> <td>70.0</td> <td>76.0</td> <td>93.0</td> </tr> <tr> <td>FasterViT-6</td> <td>53.7</td> <td>78.4</td> <td>57.1</td> <td>70.1</td> <td>76.1</td> <td>93.0</td> </tr> </table>

A、R和V2分别代表ImageNet-A、ImageNet-R和ImageNet-V2。

安装

我们提供了一个docker文件。此外，假设已安装最新版本的PyTorch，可以通过运行以下命令安装依赖项：

pip install -r requirements.txt

训练

有关所有模型的详细训练说明，请参阅TRAINING.md。

评估

可以使用以下命令在ImageNet-1K验证集上评估FasterViT模型：

python validate.py \
--model <模型名称>
--checkpoint <检查点路径>
--data_dir <imagenet路径>
--batch-size <每个GPU的批量大小>

这里--model是FasterViT变体（例如faster_vit_0_224_1k），--checkpoint是预训练模型权重的路径，--data_dir是ImageNet-1K验证集的路径，--batch-size是批量大小。我们还在这里提供了一个示例脚本。

ONNX转换

我们提供ONNX转换脚本以支持动态批量大小推理。例如，要为分辨率为576 x 960且ONNX操作集版本为17的faster_vit_0_any_res生成ONNX模型，可以使用以下命令：

python onnx_convert --model-name faster_vit_0_any_res --simplify --resolution-h 576 --resolution-w 960 --onnx-opset 17

CoreML转换

要生成FasterViT CoreML模型，请安装coremltools==5.2.0并使用我们提供的脚本。

建议使用Xcode14或更新版本对性能进行基准测试。

星标历史

第三方扩展

我们一直欢迎第三方扩展/实现和用于其他目的。以下代表其他用户的第三方贡献。

名称	链接	贡献者	框架
keras_cv_attention_models	链接	@leondgarse	Keras

如果您希望您的工作在此存储库中列出，请提出一个问题并提供给我们详细信息。

引用

如果此存储库对您的工作有用，请考虑引用FasterViT。

@article{hatamizadeh2023fastervit,
  title={FasterViT: Fast Vision Transformers with Hierarchical Attention},
  author={Hatamizadeh, Ali and Heinrich, Greg and Yin, Hongxu and Tao, Andrew and Alvarez, Jose M and Kautz, Jan and Molchanov, Pavlo},
  journal={arXiv preprint arXiv:2306.06189},
  year={2023}
}