[论文阅读笔记34] LISA (LISA: Reasoning Segmentation via Large Language Model) 代码精读

---

LISA是一个很好的Reason Segmentation的baseline, 其利用特殊的token [SEG]来微调多模态LLM和SAM的decoder来实现复杂逻辑下的prompt的推理分割. 其整体框图如下, 本篇文章精度此代码并作简单复现.

---

1. 推理流程

流程如下：

1.1 加载Tokenizer与模型

首先利用transformers库的AutoTokenizer从config文件中加载Tokenizer. LISA使用的Tokenizer是LLaMa相同的Tokenizer，填充方式是在序列右侧填充，最大长度是512:

tokenizer = AutoTokenizer.from_pretrained(
        args.version,
        cache_dir=None,
        model_max_length=args.model_max_length,  # 512
        padding_side="right",
        use_fast=False,  # 如果选择true 则模型自动进行填充 选择False的流程是首先对输入进行encode 然后再填充
    )

然后设定填充的unknown token以及本工作提出的代表分割掩码的[SEG] token:

tokenizer.pad_token = tokenizer.unk_token  # unknown token为 <unk>
args.seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]  # 在codebook中 [SEG]的id是32000
# tokenizer返回一个dict, 有两个key, input_ids表示在词汇表中的索引, attention_mask表示是否为填充的token 用于在注意力计算的时候看哪些token应该被注意

然后根据tokenizer的设定，设置model的句子开始、结束以及填充的token id:

model.config.eos_token_id = tokenizer.eos_token_id  # </s>
model.config.bos_token_id = tokenizer.bos_token_id  # <s>
model.config.pad_token_id = tokenizer.pad_token_id  # <unk>

随后，加载CLIP预训练的ViTal-large模型，作为LLaVA中的vision encoder:

model.get_model().initialize_vision_modules(model.get_model().config)
vision_tower = model.get_model().get_vision_tower()
vision_tower.to(dtype=torch_dtype)

加载CLIP的图像预处理类，包含resize、crop等，以及SAM中用到的resize类，其按照图像的最长边进行等比例resize:

clip_image_processor = CLIPImageProcessor.from_pretrained(model.config.vision_tower)
transform = ResizeLongestSide(args.image_size)

1.2 推理主要过程

按照LLaVA规定的格式先实例化一个Conversation类，这个类是自定义的一个数据类，用以保存所有的对话历史：

conv = conversation_lib.conv_templates[args.conv_type].copy()
"""
规定一开始的system设定 角色(用户和bot) 对话历史 分隔的token等
Conversation(system="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.", roles=('USER', 'ASSISTANT'), messages=(), offset=0, sep_style=<SeparatorStyle.TWO: 2>, sep=' ', sep2='</s>', version='v1', skip_next=False)
"""
conv.messages = []  # 初始化

读取文本prompt, 例如who is the oldest person? Please output segmentation mask.

读取文本后，要在文本prompt之前加入给图像预留的token <image>, 并且在<image>前后加入起止符:

prompt = DEFAULT_IMAGE_TOKEN + "\n" + prompt
prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)  # <im_start><image><im_end>\nwho is the oldest person? Please output segmentation mask.

随后组合成完整的对话prompt:

A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <im_start><image><im_end>\nwho is the oldest person? Please output segmentation mask. ASSISTANT:

读取图片，并用CLIP的image preprocess将图像缩放为224x224, 并遵循SAM的预处理将图像长边缩放至1024, 并填充至1024x1024. 两种预处理分别对应两个vision encoder.

随后将文本prompt进行tokenize, 得到input_ids:

input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt")
input_ids = input_ids.unsqueeze(0).cuda()  # [bs, length]
"""
tensor([[    1,   319, 13563,  1546,   263, 12758,  5199,   322,   385, 23116,
         21082, 20255, 29889,   450, 20255,  4076,  8444, 29892, 13173, 29892,
           322,  1248,   568,  6089,   304,   278,  5199, 29915, 29879,  5155,
         29889,  3148,  1001, 29901, 32001,  -200, 32002,  1058,   338,   278,
         23947,  2022, 29973,  3529,  1962, 10768,   362, 11105, 29889,   319,
          1799,  9047, 13566, 29901]], device='cuda:0')
"""

将input_ids输入LISA模型, 实际上走的是LLaVA, 并产生文本输出, 类似于训练时预规定的"Sure, It is [SEG]":

with torch.no_grad():
    outputs = self.generate(  # transformers库的方法 采用greedy生成 即每次选logits最大的token输出
        images=images_clip,
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        num_beams=1,
        output_hidden_states=True,  # 输出hidden state 是因为我们要取[SEG]对应的embedding来decode分割的mask
        return_dict_in_generate=True,
    )
"""
输出的句子(outputs.sequence)为:
tensor([[    1,   319, 13563,  1546,   263, 12758,  5199,   322,   385, 23116,
         21082, 20255, 29889,   450, 20255,  4076,  8444, 29892, 13173, 29892,
           322,  1248,   568,  6089,   304,   278,  5199, 29915, 29879,  5155,
         29889,  3148,  1001, 29901, 32001,  -200, 32002,  1058,   338,   278,
         23947,  2022, 29973,  3529,  1962, 10768,   362, 11105, 29889,   319,
          1799,  9047, 13566, 29901, 18585, 29892, 32000,   869,     2]],
       device='cuda:0')
翻译过来就是:
<s>A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <im_start> <im_end> who is the oldest person? Please output segmentation mask. ASSISTANT: Sure, [SEG] .</s>
"""

output_hidden_states = outputs.hidden_states[-1]  # 取最后一层的hidden state [bs, 313, 5120]
output_ids = outputs.sequences

之后是重要的步骤: 把[SEG]在输出中对应的位置取出来(得到一个mask, 只有[SEG]在的位置是1), 并在前面填充255个False, 原因是:在LLaVA的推理过程中, vision encoder会将224x224的图像切分成14x14大小的patch, 共256个, 所以等价的输出也是256个, 因此在真实的LLaVA输出中, 长度比原本多了256 - 1(<image>符).

eg_token_mask = output_ids[:, 1:] == self.seg_token_idx
# hack for IMAGE_TOKEN_INDEX (we suppose that there is only one image, and it is in the front)
seg_token_mask = torch.cat(
    [
        torch.zeros((seg_token_mask.shape[0], 255)).bool().cuda(),
        seg_token_mask,
    ],
    dim=1,
)

取出[SEG]对应的hidden state, 并经过一个MLP层用以对齐:

hidden_states.append(self.model.text_hidden_fcs[0](output_hidden_states))

last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)
pred_embeddings = last_hidden_state[seg_token_mask]

然后是SAM阶段, 对SAM预处理的image经过encoder, 然后将[SEG]对应的hidden state作为prompt, 也进行encode:

image_embeddings = self.get_visual_embs(images)  # [1, 256, 64, 64] 其中 64 = 1024 / 16 patch大小为16 x 16

multimask_output = False
pred_masks = []
for i in range(len(pred_embeddings)):
    (
        sparse_embeddings,  # [1, 1, 256]
        dense_embeddings,  # [1, 256, 64, 64]
    ) = self.model.visual_model.prompt_encoder(
        points=None,
        boxes=None,
        masks=None,
        text_embeds=pred_embeddings[i].unsqueeze(1),
    )

随后执行常规的SAM decoder过程, 得到分割掩码. 最后, 对output_id进行解码得到文本输出, 以及对mask进行可视化与保存即可.

low_res_masks, iou_predictions = self.model.visual_model.mask_decoder(
    image_embeddings=image_embeddings[i].unsqueeze(0),
    image_pe=self.model.visual_model.prompt_encoder.get_dense_pe(),  # 图像的位置编码
    sparse_prompt_embeddings=sparse_embeddings,
    dense_prompt_embeddings=dense_embeddings,
    multimask_output=multimask_output,  # False, LISA只默认分割一个物体
)
pred_mask = self.model.visual_model.postprocess_masks(  # 将低分辨率mask [224, 224]后处理回原本的分辨率
    low_res_masks,
    input_size=resize_list[i],
    original_size=original_size_list[i],
)

text_output = tokenizer.decode(output_ids, skip_special_tokens=False)

1.3 运行结果

可以看出来，它并没有很正确地分割出来最“老”的人(应该是画面最右侧的), 而是倾向于分割所有的人, 说明LISA可能存在对特定文本忽视的现象, 这在一些生成的工作中有人关注过.

2. 训练流程

2.1 数据准备与读取

LISA的训练数据包括四个任务：

• 语义分割
• 指令分割
• VQA
• 推理分割

在训练的时候, 将四种任务的若干数据集混合, 封装成统一的HybridDataset类, 在每次迭代的时候, 都随机从四个任务中挑选一个任务, 再从挑选的任务中随机选一个数据集, 再从这个数据集中随机选一个样本, 代码如下:

# HybridDataset:
def __getitem__(self, idx):
    ind = np.random.choice(list(range(len(self.datasets))), p=self.sample_rate)
    data = self.all_datasets[ind]
    inference = False
    return *data[0], inference
    
# 传到单个任务的数据集类中索引是0 但是也是随机选一个数据集之后随机选一个样本 例如对应Refferring Segment:
def __getitem__(self, idx):
    ds = random.randint(0, len(self.refer_seg_ds_list) - 1)
    ds = self.refer_seg_ds_list[ds]
    refer_seg_ds = self.refer_seg_data[ds]
    images = refer_seg_ds["images"]
    annotations = refer_seg_ds["annotations"]
    img2refs = refer_seg_ds["img2refs"]
    idx = random.randint(0, len(images) - 1)

接下来看一下每个任务的数据集是如何构建以及读取的.

2.2.1 HyBridDataset初始化

在训练的主函数(train_ds.py)中, 对混合数据集进行如下初始化:

train_dataset = HybridDataset(
    args.dataset_dir,  # 根目录 存放所有任务的所有数据集
    tokenizer,  # 采用LLaVA的tokenizer
    args.vision_tower,  # CLIP的ViT-large
    samples_per_epoch=args.batch_size  # 一个epoch的样本数 = bs * 梯度积累步数 * 一个epoch的步数 * 显卡数
    * args.grad_accumulation_steps  # 默认10 梯度累积的主要目的是在显存有限的情况下, 模拟大bs的训练效果
    * args.steps_per_epoch  # 默认500
    * world_size,
    precision=args.precision,  # 推理时候的精度 fp16/fp32
    image_size=args.image_size,  # 默认1024 SAM的输入分辨率
    num_classes_per_sample=args.num_classes_per_sample,  # 一个样本的标注中最多看几个类别 默认为3
    exclude_val=args.exclude_val,  # 是否排除验证集
    dataset=args.dataset,  # 默认四个任务都进行
    sample_rate=[float(x) for x in args.sample_rates.split(",")],  # 对每个任务的采样频率 
    # 默认语义分割: 指令分割: VQA: 因果分割 = 9: 3: 3: 1 可以看出是保住分割能力 并防止复杂prompt的过拟合
    sem_seg_data=args.sem_seg_data,  # 具体的语义分割的数据集名称
    refer_seg_data=args.refer_seg_data,  # 具体的指令分割的数据集名称
    vqa_data=args.vqa_data,  # 具体的VQA的数据集名称
    reason_seg_data=args.reason_seg_data,  # 具体的推理分割的数据集名称
    explanatory=args.explanatory,  # 这个参数是对ReasonSeg而言的, 问题是要求解释的问题("例如Please segment.. and explain why")的比例, 默认是0.1. 加入VQA数据集训练的目的也是保障模型回答问题的能力.
)

2.2.2 语义分割

在初始化函数中, 分别定义图像预处理方法, 问题模板, 回答模板以及每个数据集的样本的class名称, 图像以及label:

self.transform = ResizeLongestSide(image_size)
self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

self.short_question_list = SHORT_QUESTION_LIST  # 例如 <image> + "\n" + "Can you segment the {class_name} in this image?"
self.answer_list = ANSWER_LIST  # 例如 "It is [SEG]." "Sure, [SEG]."等

self.data2list = {}  # key: 数据集 value: (images, labels), images 和 labels为长度为N的列表, 里面存储路径
self.data2classes = {}  # key: 数据集 value: class名称的np.ndarray

# 存储每一个数据集
self.sem_seg_datas = sem_seg_data.split("||")
for ds in self.sem_seg_datas:
    classes, images, labels = eval("init_{}".format(ds))(base_image_dir)
    self.data2list[ds] = (images, labels)
    self.data2classes[ds] = classes

在__getitem__中, 首先随机选择数据集, 然后再从数据集中随机选择一个样本进行图像和label的读取和resize. 此外, 读取对应的类别(如果超过规定的数目, 就随机抽取args.num_classes_per_sample个), 这部分不再赘述.

之后, 构建问题和答案:

questions = []
answers = []
class_ids = []  # 为样本中的每一个类别创建一组问答
for sampled_cls in sampled_classes:
    text = sampled_cls

    assert len(text.split("||")) == 1
    question_template = random.choice(self.short_question_list)  # 按照模板构建问题
    questions.append(question_template.format(class_name=text.lower()))

    answers.append(random.choice(self.answer_list))  # 随机选择答案模板

    if ds in ["paco_lvis", "pascal_part"]:  # 这两个数据集是single class 特殊处理
        continue

    class_id = self.data2classes[ds].tolist().index(sampled_cls)
    class_ids.append(class_id)

# 转换为标准的prompt 即 system: A chat... Human: XXX Assistant: XXX
conversations = []
conv = conversation_lib.default_conversation.copy()

i = 0
while i < len(questions):
    conv.messages = []
    conv.append_message(conv.roles[0], questions[i])
    conv.append_message(conv.roles[1], answers[i])
    conversations.append(conv.get_prompt())
    i += 1

随后读取label中的mask, 这部分省略, 返回值是如下的格式, 其余数据集也遵循:

return (
    image_path,  # 图像路径
    image,  # 用于SAM的resize图像 应该是1024x1024
    image_clip,  # 用于CLIP的resize图像 应该是224x224
    conversations,  # 真值conversation
    masks,  # 真值masks shape: [n, h, w] n是对应的物体类别数 [h, w]是原大小
    label,  # [h, w], 原始分割标签
    resize,  # [1024, 1024]
    questions,  # 问题
    sampled_classes,  # 抽取的类别名称 list
)

2.2.3 指令分割

基本流程和语义分割是相似的, 只不过class name需要从annotation的referring中读出来:

img2refs = refer_seg_ds["img2refs"]
refs = img2refs[image_id]  # 得到图像对应的referrings

# 读取对应所有referring的文本 当然后面要根据args.num_classes_per_sample作筛选
sents = []
ann_ids = []
for ref in refs:
    for sent in ref["sentences"]:
        text = sent["sent"]
        sents.append(text)
        ann_ids.append(ref["ann_id"])
        
# 因此
sampled_classes = sampled_sents

# 后面读取mask也类似, 要根据抽取出的referring找到对应的mask

2.2.4 VQA

VQA比较特殊. 直接从数据集中读取数据即可构建conversation, 对于mask和label, 则将mask置为全0, label都置为ignore_label(255):

conversations = []
if roles[source[0]["from"]] != conv.roles[0]:
    # Skip the first one if it is not from human
    source = source[1:]
conv.messages = []
for j, sentence in enumerate(source):
    role = roles[sentence["from"]]  # 直接从数据集读取
    assert role == conv.roles[j % 2], f"{i}"
    conv.append_message(role, sentence["value"])
conversations.append(conv.get_prompt())

questions = conversations
sampled_classes = conversations

image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())

masks = torch.rand(0, *ori_size)  # 全0
label = torch.ones(ori_size) * self.ignore_label  # 全255

2.2.5 ReasonSeg

ReasonSeg最大的不同就是要处理长问话以及解释性的问话,

首先读取当前随机抽取样本的mask, 问话以及是否为一个句子:

mask, sents, is_sentence = get_mask_from_json(json_path, image)

随后看是否为需要解释的样本, 如果是的话就构建对应的问话. 其中的choice是在这种情况下, 进一步控制是解释性问题的比例, 其实是让解释性问题的占比进一步降低了.

if is_sentence:
    question_template = random.choice(self.long_question_list)
    questions.append(question_template.format(sent=text))
else:
    question_template = random.choice(self.short_question_list)
    questions.append(question_template.format(class_name=text.lower()))

# add explanation if applicable
img_name = image_path.split("/")[-1]
if self.explanatory != -1 and img_name in self.img_to_explanation:
    if choice == 0:  # [SEG] token  # 最简单的回答
        answers.append(random.choice(self.answer_list))
    elif choice == 1:  # [SEG] token + text answer  # 否则加入解释性的提问
        image_name = image_path.split("/")[-1]
        answer = self.img_to_explanation[image_name]["outputs"]
        answer = random.choice(self.answer_list) + " {}".format(answer)
        questions[-1] = (
            DEFAULT_IMAGE_TOKEN
            + "\n"
            + text
            + " {}".format(random.choice(self.explanatory_question_list))
        )
        answers.append(answer)
    elif choice == 2:  # vanilla text answer  # 不加入
        image_name = image_path.split("/")[-1]
        answer = self.img_to_explanation[image_name]["outputs"]
        questions[-1] = DEFAULT_IMAGE_TOKEN + "\n" + text
        answers.append(answer)
    else:
        raise ValueError("Not implemented yet.")
else:
    answers.append(random.choice(self.answer_list))

2.2 单步训练流程与损失计算

2.2.1 模型载入

模型载入和推理过程基本是相似的. 但是训练过程中需要用LoRA来微调LLaVA (要训练生成固定的答案模板), 具体做法如下:

lora_r = args.lora_r  # 降到秩为多少 默认为8
if lora_r > 0:
	
    # 查找模型中符合条件的线性层 并保存下来
    def find_linear_layers(model, lora_target_modules):
        cls = torch.nn.Linear
        lora_module_names = set()
        for name, module in model.named_modules():
            if (
                isinstance(module, cls)
                and all(
                    [
                        x not in name
                        for x in [
                            "visual_model",
                            "vision_tower",
                            "mm_projector",
                            "text_hidden_fcs",
                        ]
                    ]
                )
                and any([x in name for x in lora_target_modules])
            ):
                lora_module_names.add(name)
        return sorted(list(lora_module_names))

    lora_alpha = args.lora_alpha  # lora超参
    lora_dropout = args.lora_dropout  # lora超参
    lora_target_modules = find_linear_layers(
        model, args.lora_target_modules.split(",")
    )
    # 配置config 利用peft库实现lora
    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=lora_target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
    )
    # 根据lora要改变的线性层更新模型
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

2.2.2 前向传播

在每次迭代中, 首先用SAM的image encoder得到图像特征, 以及, 得到[SEG]在tokenize的输入中的位置:

image_embeddings = self.get_visual_embs(images)  # [bs, c, h, w]
batch_size = image_embeddings.shape[0]
assert batch_size == len(offset) - 1

seg_token_mask = input_ids[:, 1:] == self.seg_token_idx  # [bs, N]
seg_token_mask = torch.cat(
    [
        seg_token_mask,
        torch.zeros((seg_token_mask.shape[0], 1)).bool().cuda(),  # [bs, 1]
    ],
    dim=1,
)  # [bs, N + 1]
# hack for IMAGE_TOKEN_INDEX (we suppose that there is only one image, and it is in the front)
# 补齐255个0的理由和推理节介绍的相同
seg_token_mask = torch.cat(
    [torch.zeros((seg_token_mask.shape[0], 255)).bool().cuda(), seg_token_mask],
    dim=1,
)

然后输入LLaVA:

images_clip_list = []
for i in range(len(offset) - 1):
    start_i, end_i = offset[i], offset[i + 1]  # 该样本具有多少annotation, 即问答对的起始和最终的idx
    images_clip_i = (
        images_clip[i]
        .unsqueeze(0)
        .expand(end_i - start_i, -1, -1, -1)  # 就重复这么多遍
        .contiguous()
    )
    images_clip_list.append(images_clip_i)
images_clip = torch.cat(images_clip_list, dim=0)

output = super().forward(  # 得到LLaVA的结果
    images=images_clip,
    attention_mask=attention_masks,
    input_ids=input_ids,
    labels=labels,
    output_hidden_states=True,
)
output_hidden_states = output.hidden_states

得到最后一层的各个batch中[SEG]的embedding:

hidden_states = []

assert len(self.model.text_hidden_fcs) == 1
hidden_states.append(self.model.text_hidden_fcs[0](output_hidden_states[-1]))  # 输入FC层对齐SAM和LLavA

last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)
pred_embeddings = last_hidden_state[seg_token_mask]
seg_token_counts = seg_token_mask.int().sum(-1)  # [bs, ]

# 得到每个batch中seg token的起始位置 并获得相应的embeddings
seg_token_offset = seg_token_counts.cumsum(-1)
seg_token_offset = torch.cat(
    [torch.zeros(1).long().cuda(), seg_token_offset], dim=0
)

seg_token_offset = seg_token_offset[offset]

pred_embeddings_ = []
for i in range(len(seg_token_offset) - 1):
    start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1]
    pred_embeddings_.append(pred_embeddings[start_i:end_i])
pred_embeddings = pred_embeddings_

遍历每个embeddings, 用SAM的decoder得到mask, 这部分和推理过程相似, 不再赘述. 最后计算loss. 一个是训练LLaVA用的交叉熵损失, 用于文本输出和模板一致; 另外就是分割常用的bce和dice loss. 注意对于VQA任务的样本, mask loss理应是0.