S²-Wrapper

本仓库包含S²-Wrapper的Pytorch实现，这是一种简单的机制，能够在任何视觉模型上实现多尺度特征提取。

阅读我们的论文，了解何时在图像尺度上进行缩放比在模型规模上进行缩放更好。

何时我们不需要更大的视觉模型？
Baifeng Shi, Ziyang Wu, Maolin Mao, Xin Wang, Trevor Darrell
加州大学伯克利分校, 微软研究院

论文: https://arxiv.org/abs/2403.13043

新闻

[2024/07] 被ECCV 2024接收！
[2024/05] S²-Wrapper正式集成到NVIDIA VILA中！我们发布了带有S²的VILA-1.5-3b检查点，更多检查点正在路上！在这里查看详情。
[2024/04] S²-Wrapper正式集成到LLaVA中！我们发布了带有S²的LLaVA-1.5检查点。在这里试用。

待办事项

添加带有S²-Wrapper的LLaVA-NeXT预训练检查点。
~~添加带有S²-Wrapper的LLaVA-1.5预训练检查点。~~
~~添加对非方形图像的支持~~ 现在支持任何形状的图像。请在dev_any_shape分支中查看。该功能仍在测试中。
~~添加LLaVA与S²-Wrapper结合的示例~~

快速入门

步骤1. 通过pip安装s2wrapper。

pip install git+https://github.com/bfshi/scaling_on_scales.git

步骤2. 使用一行代码在任何视觉模型上提取多尺度特征。

假设您有一个函数（可以是model、model.forward等）接收BxCxHxW图像并输出BxNxC特征。

例如，您有一个model（如ViT-B）通过以下方式提取特征：

feature = model(x)   # 例如，x: 32*3*224*224, feature: 32*196*768

然后通过以下方式提取多尺度特征（例如，尺度为1和2）：

from s2wrapper import forward as multiscale_forward
mutliscale_feature = multiscale_forward(model, x, scales=[1, 2])   # x: 32*3*224*224, feature: 32*196*1536

上面我们假设输入是224x224的，s2wrapper会将其插值到448x448。如果原始448x448图像已经可用，我们可以通过从448x448图像插值而不是从224x224图像插值获得更好的性能。在这种情况下，通过以下方式提取224x224和448x448尺度的特征：

from s2wrapper import forward as multiscale_forward
mutliscale_feature = multiscale_forward(model, x, scales=[0.5, 1], max_split_size=224)   # x: 32*3*448*448, feature: 32*196*1536，注意我们需要设置`max_split_size=224`以将448图像分割成4个子图像。
# mutliscale_feature = multiscale_forward(model, x, img_sizes=[224, 448], max_split_size=224)   # 或者，设置`img_sizes`而不是`scales`

使用方法

s2wrapper.forward(
    model,
    input,
    scales=None,
    img_sizes=None,
    max_split_size=None,
    resize_output_to_idx=0,
    num_prefix_token=0,
    output_shape='bnc',
    split_forward=False,
)

model: 您的视觉模型或任何接收BxCxHxW图像张量并输出BxNxC特征张量的函数。

input: 输入图像张量，形状为BxCxHxW。

scales: 用于提取特征的尺度列表。例如，如果默认大小为224²，scales=[1, 2]将在224²和448²尺度上提取特征。

img_sizes: 或者，您可以为每个尺度指定图像大小，而不是指定scales。例如，对于默认大小为224²，img_sizes=[224, 448]将产生与scales=[1, 2]相同的结果。

max_split_size: 从大图像分割出的子图像的最大大小。对于每个尺度，图像将被分割成ceil(img_size_that_scale / max_split_size)**2个子图像。如果为None，则默认设置为input的大小。

resize_output_to_idx: 将最终特征图调整到哪个尺度。默认为scales或img_sizes中的第一个尺度。

num_prefix_token: 特征图中的前缀标记数量。例如，如果model返回的特征图包含1个[CLS]标记和其他空间标记，设置num_prefix_token=1。默认为0。

output_shape: 输出特征的形状。需要是bnc（例如，ViT）或bchw（例如，ConvNet）。默认为bnc。

split_forward: 是否在每个子图像上单独运行模型或将所有子图像批处理为单次运行。设置为True可以减少内存使用（在推理过程中，GPU内存使用量大致与单尺度相同）。默认为False。

示例：带有S²-Wrapper的LLaVA

S²-Wrapper已正式集成到LLaVA中（参见这里的PR）。要使用带有S²-Wrapper的LLaVA，只需安装此仓库和最新版本的LLaVA仓库，并下载下面列出的检查点。我们已发布带有S²-Wrapper的LLaVA-1.5-7B和LLaVA-1.5-13B的检查点。

模型	大小	计划	检查点	VQAv2	VizWiz	TextVQA	MMMU-val	MathVista	MM-Bench	SEED	MM-Vet
LLaVA-1.5	7B	full_ft-1e	liuhaotian/llava-v1.5-7b	78.5	50.0	58.2	36.2	25.2	64.3	65.7	31.1
LLaVA-1.5	7B	lora-1e	liuhaotian/llava-v1.5-7b-lora	79.1	47.8	58.2	-	-	66.1	-	30.2
LLaVA-1.5-S2	7B	lora-1e	bfshi/llava-v1.5-7b-s2-lora	80.0	50.1	61.0	37.7	25.3	66.2	67.9	32.4
LLaVA-1.5	13B	full_ft-1e	liuhaotian/llava-v1.5-13b	80.0	53.6	61.3	36.4	27.6	67.7	68.2	36.1
LLaVA-1.5	13B	lora-1e	liuhaotian/llava-v1.5-13b-lora	80.0	58.9	60.2	-	-	68.5	-	38.3
LLaVA-1.5-S2	13B	lora-1e	bfshi/llava-v1.5-13b-s2-lora	80.9	56.0	63.1	37.4	27.8	67.9	68.9	36.4

使用LLaVA-1.5-S2进行模型推理的示例脚本：

python3 -m llava.eval.run_llava \
    --model-path bfshi/llava-v1.5-7b-s2-lora \
    --model-base lmsys/vicuna-7b-v1.5 \
    --image-file <image> \
    --query <query> \
    --conv-mode vicuna_v1

训练。要训练带有S²-Wrapper的LLaVA，由于当前的LLaVA仓库仅支持使用S²进行评估，请额外应用这里的更改到您的LLaVA仓库，然后就可以开始了！

训练配置应与训练常规LLaVA 不使用 anyres相同（即image_aspect_ratio="resize"和mm_patch_merge_type="flat"），除了两个新的模型配置：

s2=True。这会启用S²的使用。
s2_scales="336,672,1008"。这指定了S²将在哪些图像尺度上提取特征。

示例：带有S²-Wrapper的NVIDIA VILA

S²-Wrapper已正式集成到NVIDIA VILA中。VILA是一个多模态LLM，支持多图像理解和视频理解，在多个基准测试上取得了出色的结果（例如，在MMMU上在所有开源模型中排名第一）。VILA有几种模型大小：3B、8B、13B和40B，每种还有一个量化版本（AWQ）。

目前，我们已发布了带有S²-Wrapper的VILA-3B检查点，这是您在边缘设备上运行MLLM的首选。其他模型大小的检查点正在路上！同时，欢迎在这里查看更多详情。

$~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~$	精度	VQAv2	GQA	VizWiz	SQA-I	VQA-T	POPE	MME	MMB	MMB-CN	SEED	SEED-I	MMMU (验证)	MMMU (测试)	llava-bench	MM-Vet	平均值
VILA1.5-3B	fp16	80.4	61.5	53.5	69.0	60.4	85.9	1442.44	63.4	52.7	60.9	67.9	33.3	30.8	75.9	35.4	60.2
VILA1.5-3B-S2	fp16	79.8	61.4	61.3	69.6	63.4	85.3	1431.65	62.8	52.2	60.0	66.4	32.8	31.3	76.7	38.6	60.9
VILA1.5-3B-AWQ	int4	80.0	61.1	53.8	67.8	60.4	85.9	1437.34	63.3	51.4	59.8	66.6	32.7	31.1	75.0	37.3	59.9
VILA1.5-3B-S2-AWQ	int4	79.4	61.3	62.3	69.2	63.0	85.8	1417.06	61.6	51.5	59.1	65.7	33.4	30.4	77.1	36.7	60.5

请参考VILA的原始仓库获取检查点以及关于训练、评估和部署的指导。

示例：使用S²-Wrapper的HuggingFace CLIP

使用HuggingFace CLIP视觉模型进行常规特征提取（参考：官方示例）：

from PIL import Image
import requests
from transformers import AutoProcessor, CLIPVisionModel

model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32")
processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

inputs = processor(images=image, return_tensors="pt").pixel_values

# model.forward返回一个包含"last_hidden_state"的对象，这就是我们需要的特征图
outputs = model(inputs).last_hidden_state
print(outputs.shape)  # 1*50*768

将其改为多尺度：

from PIL import Image
import requests
from transformers import AutoProcessor, CLIPVisionModel

model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32")
processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

inputs = processor(images=image, return_tensors="pt").pixel_values

# 将特征提取过程封装成一个单独的函数，该函数
# 接受图像张量作为输入并输出特征张量
def forward_features(inputs):
    return model(inputs).last_hidden_state

# 使用scales=[1, 2]提取特征。注意输出有一个[CLS]标记
# 所以设置num_prefix_token=1。
outputs = multiscale_forward(forward_feature, inputs, scales=[1, 2], num_prefix_token=1)
print(outputs.shape)  # 1*50*1536

引用

@article{shi2024we,
  title={When Do We Not Need Larger Vision Models?},
  author={Shi, Baifeng and Wu, Ziyang and Mao, Maolin and Wang, Xin and Darrell, Trevor},
  journal={arXiv preprint arXiv:2403.13043},
  year={2024}
}