Hiera：一款去繁从简的层次化视觉Transformer

这是我们ICML 2023口头报告论文的官方实现：
Hiera：一款去繁从简的层次化视觉Transformer
Chaitanya Ryali*, Yuan-Ting Hu*, Daniel Bolya*, Chen Wei, Haoqi Fan, Po-Yao Huang, Vaibhav Aggarwal, Arkabandhu Chowdhury, Omid Poursaeed, Judy Hoffman, Jitendra Malik, Yanghao Li*, Christoph Feichtenhofer*
ICML '23 口头报告 | GitHub | arXiv | BibTeX

*: 共同第一作者。

什么是Hiera？

Hiera 是一种快速、强大且最重要的是简单的层次化视觉Transformer。它在广泛的图像和视频任务中均优于当前最先进水平，同时速度也快得多。

它是如何工作的？

像ViT这样的视觉Transformer在整个网络中使用相同的空间分辨率和特征数量。但这样做效率不高：早期层并不需要那么多特征，而后期层也不需要那么高的空间分辨率。先前的层次化模型，如ResNet，通过在开始时使用较少的特征、在结束时使用较低的空间分辨率来解决这一问题。

已经出现了一些采用这种层次化设计的特定领域视觉Transformer，例如Swin或MViT。然而，在追求ImageNet-1K上完全监督训练下的最先进结果的过程中，这些模型为了弥补ViT缺乏的空间偏差，不断增加专门模块，导致模型越来越复杂。尽管这些改进确实产生了效果良好、浮点运算量诱人的模型，但从底层来看，增加的复杂性反而使这些模型整体变得更慢。

我们证明，许多这样的冗余部分其实是不必要的。与其通过架构调整手动添加空间偏置，不如直接让模型自己学习这些偏置。通过使用MAE进行训练，我们可以简化或移除现有Transformer中的所有这些复杂模块，并在此过程中提高精度。最终得到的就是Hiera——一种极其高效且简单的架构，它在多项图像和视频识别任务中都超越了当前最先进水平。

最新消息

• [2024.03.02] 代码许可证已放宽至更宽松的Apache 2.0！模型许可证保持不变。
• [2023.06.12] 增加了更多in1k模型和一些视频示例，请参阅inference.ipynb（v0.1.1）。
• [2023.06.01] 初始发布。

更多详情请参阅变更日志。

安装

Hiera需要一个相对较新的torch版本。之后，您可以通过pip安装Hiera：

pip install hiera-transformer

本仓库应支持最新的timm版本，但timm是一个不断更新的包。如果您在使用较新版本的timm时遇到问题，请提交issue。

从源码安装

如果使用torch hub，则无需安装hiera包。但是，如果您想基于Hiera进行开发，从源码安装可能是个不错的选择：

git clone https://github.com/facebookresearch/hiera.git
cd hiera
python setup.py build develop

模型库

请注意，模型权重与代码采用不同的许可证。更多详情请参阅模型许可证。

Torch Hub

在这里，我们提供了Hiera的模型检查点。列出的每个模型都可以在torch hub上访问，即使未安装hiera-transformer包亦可，例如以下代码会初始化一个在ImageNet-1K上预训练并微调的基础模型：

model = torch.hub.load("facebookresearch/hiera", model="hiera_base_224", pretrained=True, checkpoint="mae_in1k_ft_in1k")

如果您只需要带有MAE预训练的模型，可以将checkpoint替换为"mae_in1k"。此外，如果您还想加载MAE解码器（例如继续预训练），可以在模型名称前加上mae_，例如：

model = torch.hub.load("facebookresearch/hiera", model="mae_hiera_base_224", pretrained=True, checkpoint="mae_in1k")

注意： 我们的MAE模型是使用归一化的像素损失训练的。这意味着在模型预测之前，补丁已经被归一化了。如果您想可视化预测结果，就需要使用可见的补丁对其进行反归一化处理（这可能会奏效，但并不完美），或者直接使用真实标签进行反归一化。更多模型名称及对应的checkpoint名称见下文。

Hugging Face Hub

此仓库也支持 🤗 hub。在安装了 hiera-transformer 和 huggingface-hub 包后，您可以简单地运行如下代码来加载模型：

from hiera import Hiera
model = Hiera.from_pretrained("facebook/hiera_base_224.mae_in1k_ft_in1k")  # 先用 MAE 预训练再在 ImageNet-1K 上微调的模型
model = Hiera.from_pretrained("facebook/hiera_base_224.mae_in1k") # 仅 MAE 预训练，未微调

请使用下方模型库中的 <model_name>.<checkpoint_name>。

如果您想保存模型，可以使用 model.config 作为配置文件，例如：

model.save_pretrained("hiera-base-224", config=model.config)

图像模型

模型	模型名称	预训练模型 (IN-1K MAE)	微调模型 (IN-1K 监督学习)	IN-1K Top-1 (%)	A100 fp16 速度 (im/s)
Hiera-T	hiera_tiny_224	mae_in1k	mae_in1k_ft_in1k	82.8	2758
Hiera-S	hiera_small_224	mae_in1k	mae_in1k_ft_in1k	83.8	2211
Hiera-B	hiera_base_224	mae_in1k	mae_in1k_ft_in1k	84.5	1556
Hiera-B+	hiera_base_plus_224	mae_in1k	mae_in1k_ft_in1k	85.2	1247
Hiera-L	hiera_large_224	mae_in1k	mae_in1k_ft_in1k	86.1	531
Hiera-H	hiera_huge_224	mae_in1k	mae_in1k_ft_in1k	86.9	274

每个模型的输入为 224×224 的图像。

视频模型

模型	模型名称	预训练模型 (K400 MAE)	微调模型 (K400)	K400 (3×5 视图) Top-1 (%)	A100 fp16 速度 (clip/s)
Hiera-B	hiera_base_16x224	mae_k400	mae_k400_ft_k400	84.0	133.6
Hiera-B+	hiera_base_plus_16x224	mae_k400	mae_k400_ft_k400	85.0	84.1
Hiera-L	hiera_large_16x224	mae_k400	mae_k400_ft_k400	87.3	40.8
Hiera-H	hiera_huge_16x224	mae_k400	mae_k400_ft_k400	87.8	20.9

每个模型的输入为 16 帧 224×224 的视频帧，时间步长为 4。

注意： 此处列出的速度是在 未使用 PyTorch 优化的缩放点积注意力机制 的情况下测得的。如果您使用的是 PyTorch 2.0 或更高版本，您的推理速度可能会比这里列出的速度更快。

使用方法

本仓库实现了运行 Hiera 模型进行推理的代码。该仓库仍在开发中。目前我们已有的功能和未来计划如下：

• 图像推理
  • 实现 MAE
• 视频推理
  • 实现 MAE
• 完整的模型库
• 训练脚本

有关如何使用 Hiera 的示例，请参阅 examples。

推理

有关如何准备推理数据的示例，请参阅 examples/inference。

您可以通过 torch hub 或 🤗 hub 实例化模型，也可以通过 从源码安装并运行以下代码来实例化模型：

import hiera
model = hiera.hiera_base_224(pretrained=True, checkpoint="mae_in1k_ft_in1k")

然后您可以像其他模型一样进行推理：

output = model(x)

视频推理的方式相同，只需使用 16x224 的模型即可。

注意：为了提高效率，Hiera 在网络开始时会对 token 进行重新排序（参见 hiera_utils.py 中的 Roll 和 Unroll 模块）。因此，默认情况下，token 并不是按空间顺序排列的。如果您希望将中间特征图用于下游任务，可以在运行模型时传递 return_intermediates 标志：

output, intermediates = model(x, return_intermediates=True)

MAE 推理

默认情况下，这些模型不包含 MAE 解码器。如果您希望使用解码器或计算 MAE 损失，可以通过以下方式实例化一个带有 MAE 的版本：

import hiera
model = hiera.mae_hiera_base_224(pretrained=True, checkpoint="mae_in1k")

当您对模型进行推理时，它将返回一个四元组 (loss, predictions, labels, mask)，其中 predictions 和 labels 仅对应于 被删除的 tokens。返回的 mask 将在 token 可见时为 True，而在 token 被删除时为 False。您可以通过在推理时传入 mask_ratio 来更改掩码比例：

loss, preds, labels, mask = model(x, mask_ratio=0.6)

对于图像，默认的掩码比例是 0.6，但对于视频则应设置为 0.9。详情请参阅论文。

注意： 我们在 MAE 预训练中使用 归一化的像素目标，这意味着在模型预测之前，每个 patch 都会单独进行归一化处理。因此，在可视化之前，您需要使用真实值对它们进行反归一化处理。详细信息请参阅 hiera_mae.py 中的 get_pixel_label_2d 函数。

基准测试

我们提供了一个便于基准测试的脚本。请参阅 examples/benchmark，了解如何使用它。

缩放点积注意力

PyTorch 2.0 引入了优化的缩放点积注意力，它可以显著加速 Transformer 模型。我们在最初的基准测试中并未使用这一功能，但由于这是一种无需额外工作即可实现的性能提升，因此本仓库会在可用时自动启用它。要享受这一优势，请确保您的 PyTorch 版本为 2.0 或更高。

训练

即将推出。

引用

如果您在工作中使用了 Hiera 或此代码，请引用以下文献：

@article{ryali2023hiera,
  title={Hiera: A Hierarchical Vision Transformer without the Bells-and-Whistles},
  author={Ryali, Chaitanya and Hu, Yuan-Ting and Bolya, Daniel and Wei, Chen and Fan, Haoqi and Huang, Po-Yao and Aggarwal, Vaibhav and Chowdhury, Arkabandhu and Poursaeed, Omid and Hoffman, Judy and Malik, Jitendra and Li, Yanghao and Feichtenhofer, Christoph},
  journal={ICML},
  year={2023}
}

许可证

本项目的代码采用Apache License, Version 2.0 许可，而模型权重则采用知识共享署名-非商业性使用 4.0 国际许可协议 许可。

有关代码许可证的更多详细信息，请参阅 LICENSE；有关模型权重许可证的更多详细信息，请参阅 LICENSE.models。

贡献

请参阅 contributing 和 行为准则。

Hiera 快速上手指南

Hiera 是一个分层视觉 Transformer（Hierarchical Vision Transformer），以其简单、快速和强大的性能著称。它在多种图像和视频任务上超越了现有技术，同时保持了极高的推理速度。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 3.8 或更高版本
• PyTorch: 建议安装较新版本的 PyTorch (支持 Torch Hub)
• 依赖库: timm (Hiera 依赖此库，若遇到版本兼容问题请提 Issue)

提示：如果您在中国大陆地区，建议使用国内镜像源加速 Python 包和模型权重的下载。

• Pip 镜像：-i https://pypi.tuna.tsinghua.edu.cn/simple
• Hugging Face 镜像：设置环境变量 HF_ENDPOINT=https://hf-mirror.com

安装步骤

您可以选择通过 pip 直接安装，或者从源码安装以便进行开发。

方式一：通过 Pip 安装（推荐）

这是最简单的使用方式，适用于直接加载模型进行推理。

pip install hiera-transformer

国内用户加速命令：

pip install hiera-transformer -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果您需要修改代码或进行二次开发，建议克隆仓库并从源码安装。

git clone https://github.com/facebookresearch/hiera.git
cd hiera
python setup.py build develop

注意：如果使用 Torch Hub 加载模型，其实不需要安装 hiera-transformer 包，但为了开发便利性，源码安装是更好的选择。

基本使用

Hiera 支持通过 Torch Hub 或 Hugging Face Hub 直接加载预训练模型，无需手动下载权重文件。以下是最简单的图像分类模型加载与初始化示例。

方法 A：使用 Torch Hub (无需安装 hiera 包)

即使未安装 hiera-transformer 包，也可直接通过 Torch Hub 加载模型。

import torch

# 加载在 ImageNet-1k 上预训练并微调过的 Base 模型
model = torch.hub.load("facebookresearch/hiera", model="hiera_base_224", pretrained=True, checkpoint="mae_in1k_ft_in1k")

# 如果只需要 MAE 预训练权重（未微调），可将 checkpoint 改为 "mae_in1k"
# model = torch.hub.load("facebookresearch/hiera", model="hiera_base_224", pretrained=True, checkpoint="mae_in1k")

model.eval()

方法 B：使用 Hugging Face Hub (需安装 hiera-transformer)

如果您已安装 hiera-transformer 和 huggingface-hub，可以使用更简洁的 API。

from hiera import Hiera

# 加载模型 (格式：<模型名称>.<检查点名称>)
# 加载经过 MAE 预训练并在 ImageNet-1k 微调的模型
model = Hiera.from_pretrained("facebook/hiera_base_224.mae_in1k_ft_in1k")

# 或者仅加载 MAE 预训练模型
# model = Hiera.from_pretrained("facebook/hiera_base_224.mae_in1k")

model.eval()

方法 C：本地导入 (需安装 hiera-transformer)

如果您通过 pip 或源码安装了包，也可以直接导入：

import hiera

# 初始化模型
model = hiera.hiera_base_224(pretrained=True, checkpoint="mae_in1k_ft_in1k")
model.eval()

可用模型速查

模型规格	模型名称 (model)	ImageNet-1K Top-1 准确率	备注
Tiny	hiera_tiny_224	82.8%	速度最快
Small	hiera_small_224	83.8%
Base	hiera_base_224	84.5%	推荐默认使用
Base+	hiera_base_plus_224	85.2%
Large	hiera_large_224	86.1%
Huge	hiera_huge_224	86.9%	精度最高

注：所有图像模型输入分辨率均为 224x224。视频模型名称类似，但包含帧数信息（如 hiera_base_16x224）。