PixelLM

任中伟*、黄志成*、魏云超†、赵耀、傅冬梅、冯佳时和金晓杰*†‡

* 共同第一作者，†通讯作者，‡项目负责人

北京交通大学、北京科技大学、字节跳动、鹏城实验室

:fire: 新闻

[2024.02.27] 我们的PixelLM已被CVPR 2024接收！
[2023.12] 我们发布了代码和数据集。

亮点

我们提出了PixelLM，这是一种新型的用于像素级推理和理解的大型多模态模型（LMM）。PixelLM能够熟练处理具有任意数量开放集目标和不同推理复杂度的任务。其设计保留了LMM的基本结构，同时避免了额外昂贵的分割模型，提高了效率并增强了对各种应用的迁移能力。
我们构建了MUSE，一个高质量的多目标推理分割数据集，为未来的研究提供了模型训练和评估支持。利用GPT-4V辅助的数据策划流程，我们创建了24.6万个问答对，涵盖了90万个实例。我们广泛的消融研究证实了该数据集在激发模型像素推理能力方面的有效性。
PixelLM在一系列基准测试中取得了新的最先进结果，显著超越了竞争方法。

简介

尽管大型多模态模型（LMM）取得了显著进展，但在涉及多个开放世界目标的图像推理任务中生成像素级掩码仍然是一个挑战。为了弥补这一差距，我们引入了PixelLM，这是一个用于像素级推理和理解的有效且高效的LMM。PixelLM的核心是一个新颖的轻量级像素解码器和一个全面的分割码本。解码器从编码详细目标相关信息的码本标记的隐藏嵌入中高效生成掩码。通过这种设计，PixelLM与流行的LMM结构保持一致，避免了对额外昂贵的分割模型的需求。此外，我们提出了一种目标细化损失，以增强模型区分多个目标的能力，从而大大提高了掩码质量。为了推进该领域的研究，我们构建了MUSE，一个高质量的多目标推理分割基准。PixelLM在各种像素级图像推理和理解任务中表现出色，在多个基准测试中超越了已建立的方法，包括MUSE、单一和多重指代分割。全面的消融实验证实了每个提出组件的有效性。

视频

架构

PixelLM具有简化的架构，包括四个主要部分：i) 预训练的CLIP-ViT视觉编码器，与文本对齐；ii) 大型语言模型；iii) 轻量级像素解码器；iv) 分割码本。PixelLM处理图像和查询文本，为不同目标生成交错的文本描述和相应的掩码。PixelLM的核心是新颖的轻量级解码器和全面的分割码本。码本包含可学习的标记，编码与不同视觉尺度上目标引用相关的上下文和知识。然后，像素解码器基于来自码本标记的隐藏嵌入和图像特征生成目标掩码。得益于这种设计，PixelLM可以生成高质量的掩码，无需外部分割模型，显著提高了效率。此外，我们提出了一种目标细化损失，以增强模型区分多个目标的能力，从而进一步提高掩码质量。

MUSE数据集

为了促进该研究领域的模型训练和评估，我们开发了MUSE，这是第一个全面的多目标推理分割数据集。MUSE以其开放集概念、详细的对象描述、复杂的多目标问答对和实例级掩码标注而突出。具体来说，我们将图像中所有实例类别名称和相应的边界框坐标输入GPT-4V。通过精心设计的提示，GPT-4V自主选择实例构建与图像内容相关的问答对。上图左侧面板展示了我们GPT-4V数据生成流程中使用的提示。右侧面板展示了生成数据的示例。

数据集可以从这个链接下载

安装

pip install -r requirements.txt
pip install flash-attn --no-build-isolation

训练

训练数据准备

除了我们的MUSE数据外，其他训练数据和数据文件的组成方式与LISA相同。您还应该在refer_seg路径下添加COCO train2017和COCO val 2017 从上述链接下载MUSE数据并按以下方式组织：

├── dataset
│   ├── ade20k
│   │   ├── annotations
│   │   └── images
│   ├── coco
│   │   └── train2017
│   │       ├── 000000000009.jpg
│   │       └── ...
│   ├── cocostuff
│   │   └── train2017
│   │       ├── 000000000009.png
│   │       └── ...
│   ├── llava_dataset
│   │   └── llava_instruct_150k.json
│   ├── mapillary
│   │   ├── config_v2.0.json
│   │   ├── testing
│   │   ├── training
│   │   └── validation
│   ├── reason_seg
│   │   └── ReasonSeg
│   │       ├── train
│   │       ├── val
│   │       └── explanatory
│   ├── refer_seg
│   │   ├── images
│   │   |   ├── saiapr_tc-12 
│   │   |   └── mscoco
│   │   |       └── images
│   │   |           ├── train2014
│   │   |           ├── train2017
│   │   |           └── val2017
│   │   ├── refclef
│   │   ├── refcoco
│   │   ├── refcoco+
│   │   └── refcocog
│   ├── vlpart
│   │   ├── paco
│   │   │   └── annotations
│   │   └── pascal_part
│   │       ├── train.json
│   │       └── VOCdevkit
│   └── muse
│       ├── muse_train.json
│       ├── muse_val.json
│       ├── muse_test_less.json
│       └── muse_test_many.json

预训练权重

LLaVA

我们的训练过程首先需要加载LLaVA的预训练权重。对于PixelLM-7B，我们使用LLaVA-Lightning-7B-v1-1，对于PixelLM-13B，我们使用liuhaotian/llava-llama-2-13b-chat-lightning-preview。

训练

deepspeed --master_port=24999 train_ds.py \
  --version="PATH_TO_LLaVA" \
  --dataset_dir='./dataset' \
  --dataset="sem_seg||refer_seg||vqa||multi_reason_seg" \
  --sample_rates="2,9,2,4" \
  --exp_name="pixellm-7b" \
  --vision-tower='openai/clip-vit-large-patch14-336' \
  --seg_token_num=3 \
  --num_classes_per_question=3 \
  --batch_size=2 \
  --pad_train_clip_images \
  --preprocessor_config='./configs/preprocessor_448.json' \
  --resize_vision_tower \
  --resize_vision_tower_size=448 \
  --vision_tower_for_mask \
  --use_expand_question_list \
  --image_feature_scale_num=2 \
  --separate_mm_projector \

训练完成后，获取完整模型权重：

cd ./runs/pixellm-7b/ckpt_model && python zero_to_fp32.py . ../pytorch_model.bin

合并LoRA权重

合并pytorch_model.bin的LoRA权重，将生成的模型以Hugging Face格式保存到您指定的路径：

CUDA_VISIBLE_DEVICES="" python merge_lora_weights_and_save_hf_model.py \
  --version="PATH_TO_LLaVA" \
  --weight="PATH_TO_pytorch_model.bin" \
  --save_path="PATH_TO_SAVED_MODEL"

推理

使用chat.py与PixelLM-7B或PixelLM-13B进行对话

python3 ./chat.py --version="./runs/PixeLM-13B/hf_model" --precision='bf16' --seg_token_num=3  \
--pad_train_clip_images \
--preprocessor_config='./configs/preprocessor_448.json' \
 --resize_vision_tower \
--resize_vision_tower_size=448 \
--vision-tower='openai/clip-vit-large-patch14-336' \
--vision_tower_for_mask \
--image_feature_scale_num=2 \
--separate_mm_projector

数据集

json文件中的每个字典包含以下键：

- "questions"：文本问题。
- "text_answers"：问题的文本答案，其中涉及的对象实例后跟{seg}符号。
- "answers"：每个答案中包含的实例注释。

少量问题和答案对不包含text_answer键。相反，文本答案由答案中每个对象实例的描述组成。数据加载文件将根据不同情况加载数据。

引用

如果您发现这个项目对您的研究有用，请考虑引用：

@article{ren2023pixellm,
  title={PixelLM: Pixel Reasoning with Large Multimodal Model},
  author={Zhongwei Ren, Zhicheng Huang, Yunchao Wei, Yao Zhao, Dongmei Fu, Jiashi Feng, Xiaojie Jin},
  journal={arXiv preprint arXiv:2312.02228},
  year={2023}
}