[CVPR 2024] AM-RADIO：聚合视觉基础模型 - 将所有领域归一化

[CVPR 2024] AM-RADIO：聚合视觉基础模型 - 将所有领域归一化的官方PyTorch实现。

Mike Ranzinger、Greg Heinrich、Jan Kautz、Pavlo Molchanov。

英伟达研究院

如有商业合作，请访问我们的网站并提交表单：英伟达研究院授权

[论文][引用]

新闻/发布

• [2024.7.22] 🔥 RADIOv2.5 ViT-B/16和ViT-L/16已发布。对于VLLM任务，RADIOv2.5-B与RADIOv2相当或更好，而RADIOv2.5-L则明显更优！详见技术报告。
• [2024.4.30] 🔥 README更新了更多指标，arXiv更新了新结果。
• [2024.3.21] 🔥 RADIOv2.1发布。使用bf16训练，提升了指标！
• [2024.2.26] AM-RADIO论文已被CVPR 2024接收
• [2024.2.15] RADIOv2发布。使用DFN CLIP、OpenAI CLIP、DINOv2、SAM教师模型进行训练。注意之前的模型未使用SAM教师模型。
• [2024.1.5] 初始GitHub仓库发布。

摘要

AM-RADIO是一个将大型视觉基础模型蒸馏成单一模型的框架。 RADIO，一种新的视觉基础模型，在各种视觉领域表现卓越，可作为视觉骨干网络的优秀替代品。它通过蒸馏整合了CLIP变体、DINOv2和SAM，保留了文本定位和分割对应等独特特性。在ImageNet零样本（+6.8%）、kNN（+2.39%）和线性探测分割（+3.8%）以及视觉语言模型（LLaVa 1.5提升至1.5%）方面优于教师模型，可扩展到任意分辨率，支持非方形图像。我们还提供了一个高效变体E-RADIO，比CLIP和DINOv2快6-10倍。

快速开始和模型版本：

最新的模型版本是RADIOv2。一旦有新模型可用，我们将更新描述。

TorchHub

要在TorchHub中加载，请使用以下命令：

import torch
model_version="radio_v2.5-l" # 用于RADIOv2.5-L模型（ViT-L/16）
#model_version="radio_v2.5-b" # 用于RADIOv2.5-B模型（ViT-B/16）
#model_version="e-radio_v2" # 用于E-RADIO
model = torch.hub.load('NVlabs/RADIO', 'radio_model', version=model_version, progress=True, skip_validation=True)
model.cuda().eval()
x = torch.rand(1, 3, 512, 512, device='cuda')

if "e-radio" in model_version:
    model.model.set_optimal_window_size(x.shape[2:]) #其中期望输入图像的(height, width)元组。

# RADIO期望输入值在[0, 1]之间。它会自动将其归一化为均值0标准差1。
summary, spatial_features = model(x)

# RADIO还支持混合精度运行：
with torch.autocast('cuda', dtype=torch.bfloat16):
    summary, spatial_features = model(x)

# 如果你更愿意预先归一化输入，可以这样做：
conditioner = model.make_preprocessor_external()

# 现在，模型不会改变输入，用户需要在调用`model(cond_x)`之前调用`cond_x = conditioner(x)`。
# 如果你想将条件处理移到现有的数据处理流程中，你很可能会这样做。
with torch.autocast('cuda', dtype=torch.bfloat16):
    cond_x = conditioner(x)
    summary, spatial_features = model(cond_x)

对于之前的版本，使用radio_v1、radio_v2、radio_v2.1，或eradio_v1用于E-RADIO模型。

请参阅快速开始部分了解更多使用详情。有关如何加载适配器（特定教师头）的信息也可在快速开始部分找到。

结果

模型统计和总结指标：

对于总结结果，我们使用模型的总结令牌。对于零样本，我们对大多数模型使用相应的语言嵌入。对于RADIO模型，我们使用来自DFN CLIP 378模型的语言嵌入。

分割指标：

• 分割设置：线性探测，简单头部
• 对于SAM COCO结果，我们用相应的RADIO模型替换SAM模型的视觉主干网络。解码器从原始模型中冻结。

LLaVa 1.5中的视觉-语言模型性能指标：

我们替换视觉主干网络并保持与LLaVa 1.5相同的LLM和训练方法：

*注意：我们在432px分辨率下运行RADIOv2.1，在768px分辨率下运行RADIOv2.5-X。虽然这看起来可能不公平，但实际上是因为 RADIOv2.1中的模式切换问题阻止了它在高于432的分辨率上获得强大的结果。从RADIOv2.5开始， 我们已经修复了模式切换问题，这使我们能够增加输入分辨率，从而在各项指标上都取得了显著的改进。 详细信息请参见技术报告。

探测3D感知能力

探测3D感知能力：我们使用来自探测视觉基础模型的3D感知能力的代码，并 评估我们的RADIO模型及其教师模型在单目深度、 表面法线和多视图对应任务上的表现，使用 NAVI数据集。对于每个任务，我们报告平均 所有阈值的准确率。RADIO保留了DINOv2的特征，并且 表现比CLIP类似模型好得多。

详细用法

import torch

# 如果不提供`version`参数，将返回最新的ViT版本。
model = torch.hub.load('NVlabs/RADIO', 'radio_model', version='radio_v2.5-l', progress=True)
model.cuda().eval()

x = torch.rand(1, 3, 224, 224, device='cuda')

# 注意：RADIO模型期望输入的值在[0, 1]范围内
# 注意2：`radio_v1`是一个ViT-H/14模型，支持输入尺寸范围为`224 < dim < 1008`
#           其中每个维度必须能被14整除。
#           支持非方形输入。
# 注意3：`radio_v2`是一个ViT-H/16模型，支持输入尺寸范围为`224 < dim < 2048`
#           其中每个维度必须能被16整除。
summary, spatial_features = model(x)

# RADIO还支持混合精度运行，如下所示：
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
    summary, spatial_features = model(x)

HuggingFace

from PIL import Image
from transformers import AutoModel, CLIPImageProcessor

hf_repo = "nvidia/RADIO" # 对于RADIO。
# hf_repo = "nvidia/E-RADIO" # 对于E-RADIO。

image_processor = CLIPImageProcessor.from_pretrained(hf_repo)
model = AutoModel.from_pretrained(hf_repo, trust_remote_code=True)
model.eval().cuda()

image = Image.open('./examples/image1.png').convert('RGB')
pixel_values = image_processor(images=image, return_tensors='pt').pixel_values
pixel_values = pixel_values.to(torch.bfloat16).cuda()

summary, features = model(pixel_values)

使用方法

RADIO和E-RADIO将返回一个包含两个张量的元组。 summary类似于ViT中的cls_token，旨在表示整个图像的一般概念。 它的形状为$(B,C)$，其中$B$是批次维度，$C$是某个通道数。 spatial_features表示更局部化的内容，适用于密集任务（如语义分割）或集成到LLM中。 RADIO和E-RADIO返回不同形状的空间特征：

• RADIO：空间特征的形状为$(B,T,D)$，其中$T$是展平的空间标记，$D$是空间特征的通道数。注意，通常$C \neq D$。
• E-RADIO：空间特征的形状为$(B,H,W,D)$，其中$H$是高度，$W$是空间特征的宽度。 对于RADIO，可以使用模型的下采样大小结合输入张量形状来转换为空间张量格式。对于'radio_v1'，patch大小为14。

from einops import rearrange
spatial_features = rearrange(spatial_features, 'b (h w) d -> b d h w', h=x.shape[-2] // patch_size, w=x.shape[-1] // patch_size)

得到的张量形状将为(B,D,H,W)，这是计算机视觉模型中常见的形状。

RADIOv1/v2 注意事项

我们训练这个模型以适应灵活的输入维度。它支持任意输入大小。返回的模型设置了一些有用的属性，您可以查询：

model.patch_size: int
model.max_resolution: int # (图像在任何维度上都不能大于此值)
model.preferred_resolution: Tuple[height, width] # 这是RADIO主要训练的分辨率，对于摘要任务可能会产生最佳结果。
                                                 # 密集任务需要实验来找到最佳分辨率。
model.window_size: Optional[int] # 如果指定了`vitdet_window_size`，这就是该值
model.min_resolution_step: int # 结合`patch_size`和`window_size`来定义图像每个维度必须是多少的倍数。
                               # 例如，如果`patch_size == 16`，那么宽度和高度必须是16的倍数
                               # 如果`patch_size == 14`且`window_size == 8`，那么宽度和高度必须是14*8的倍数

# 为方便起见，您也可以调用这个函数来获取给定图像的最近有效输入大小
nearest_height, nearest_width = model.get_nearest_supported_resolution(height=1024, width=1024)

RADIO允许非方形输入。实际上，RADIOv1和RADIOv2在允许较大图像维度变化，只固定较小维度时，都能获得更高的零样本分类分数。

适配器

(目前仅支持TorchHub)

您可以额外指定模型适配器以实现额外的行为。目前，'clip'是唯一支持的适配器。在这种模式下，radio将返回一个包含元组的字典：

model = torch.hub.load(..., adaptor_names='clip', ...)

output = model(x)

bb_summary, bb_features = output['backbone']
clip_summary, clip_features = output['clip']

参考examples/zero_shot_imagenet.py以获取使用示例。

预处理

默认情况下，RADIO期望输入图像的归一化值在[0, 1]范围内。如果您已经有现有的数据管道，并且希望在那里进行条件处理而不是在RADIO模型内部，您可以调用此函数：

preprocessor = model.make_preprocessor_external()

images = preprocessor(images)
...
output = model(images)

E-RADIO限制

E-RADIO是RADIO的一个更高效的变体，但它有一些限制：

• E-RADIO原生仅支持大小可被32整除的图像。支持其他分辨率，但可能会导致性能下降。
• E-RADIO的性能对第3和第4块中窗口注意力的窗口大小敏感。为获得最佳性能，请自动调整输入分辨率的窗口大小：model.model.set_optimal_window_size(IMAGE_SHAPE)，其中IMAGE_SHAPE是输入图像的(高度, 宽度)元组。

中间层激活

(目前仅支持RADIO)

在推理过程中可以使用forward_intermediates()方法获取中间层激活。 示例：

outputs = model.forward_intermediates(images, indices=[7, 15, 23, 31])

训练

即将推出

Star历史

[此处显示Star历史图表]

引用RADIO

如果您发现这个仓库有用，请考虑给予星标和引用：

CVPR 2024 引用：

@InProceedings{Ranzinger_2024_CVPR,
    author    = {Ranzinger, Mike and Heinrich, Greg and Kautz, Jan and Molchanov, Pavlo},
    title     = {AM-RADIO: Agglomerative Vision Foundation Model Reduce All Domains Into One},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2024},
    pages     = {12490-12500}
}

ArXiv 引用：

@misc{ranzinger2023amradio,
      title={AM-RADIO: Agglomerative Model -- Reduce All Domains Into One},
      author={Mike Ranzinger and Greg Heinrich and Jan Kautz and Pavlo Molchanov},
      year={2023},
      eprint={2312.06709},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

许可证

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

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