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 训练代码
• ImageNet-1K 预训练模型
• 任意分辨率 FasterViT
• FasterViT pip 包发布
• 添加从 ImageNet 预训练权重初始化任意分辨率 FasterViT 的能力
• ImageNet-21K 预训练模型
• 检测代码 + 模型

结果 + 预训练模型

ImageNet-1K

FasterViT ImageNet-1K 预训练模型

ImageNet-21K

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

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

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

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

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或更新版本对性能进行基准测试。

星标历史

第三方扩展

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

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

引用

如果此存储库对您的工作有用，请考虑引用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}
}

许可证

版权所有 © 2023，NVIDIA Corporation。保留所有权利。

本作品根据NVIDIA Source Code License-NC提供。点击此处查看此许可证的副本。

有关timm存储库的许可信息，请参阅其存储库。

有关ImageNet数据集的许可信息，请参阅ImageNet官方网站。

致谢

此存储库基于timm存储库构建。我们感谢Ross Wrightman创建和维护这个高质量的库。