LISA后探索

本文接着以下这篇文档:

LISA图像分割模型的上手和微调

本文记录使用个人小数据集实现微调、评估、预测LISA模型之后,进行的更细致的探索,会使用LISA本身用到的大型数据集。

🔥概述

  • 模型:LISA;

  • 任务:各类数据集针对的任务主要有语义分割(这里暂且包含实例分割)、推理分割、指代分割(短语)、视觉问答VQA。其中语义分割的数据集是纯视觉(就给图像和掩码标注等),后三个则带有文本。对于LISA而言即便语义分割也会添加prompt,训练时都是一样要输入图文得到掩码;

  • 数据集:LISA用的数据集就已经很多了,可以从中选择,暂时先考虑稍微简单点的语义分割和指代分割,其文件结构如下:

├── 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
│   │   ├── refclef
│   │   ├── refcoco
│   │   ├── refcoco+
│   │   └── refcocog
│   └── vlpart
│       ├── paco
│       │   └── annotations
│       └── pascal_part
│           ├── train.json
│           └── VOCdevkit

调研了一下其中的部分数据集,都是比较有代表性和热度的,主要就是ade20k。

数据集 任务 规模(train/val) 排行榜 描述
ade20k 语义分割 20k/2k 验证集全数据集 这个数据集的榜有比较多模型参与,大小很合适(20k、2k),适合先行调试用。
COCO
COCO-stuff
语义分割 118k/5k COCO-stuff 前者只有图像,后者只有掩码标注,该数据集排行参与者也不多(可能主要在目标检测多)。
MSCOCO 语义分割 83k/40.5k MSCOCO 只有图,最好的就是HyperSeg,断档高(但是有用额外数据集训练),这个榜没多少模型。
refCOCO
refCOCO+
refCOCOg
指代分割 (短语/目标/图像)
142k/50k/20k、
142k/50k/20k、
95k/50k/26k
refCOCOrefCOCO+refCOCOg 为MSCOCO数据集做文本标注,多模态的分割模型都有用这个模型。
  • 指标:一般都是mIoU(或者叫gIoU)和cIoU。

🔥技术原理

LISA网络结构

这是通过阅读源码得到的更完整的LISA的结构,可以看到LISA使用SAM(用的第一代)时,提示编码器不再像SAM原作一样把点、框、掩码作为提示,而是把多模态大模型输出的词元经过MLP编码后作为提示。

🚀数据流动

LISA数据流动示意

使用batch size为1的图片进行输入,他可以得到如上图所示的数据变化,设置每张图随机选择其中3个mask作为预测目标。

  1. 左边图像经过裁切、复制操作得到(3目标,3通道,224w,224h)的张量,输入视觉编码器得到(3目标,256长度,1024通道),再映射得到4096通道,此时它就是相当于是文本的token了。

  2. 而文本自然是对3个mask的prompt指令,分词编码后长度为41(不固定),图文编码合并后输入llama(正儿八经的LLM);

  3. llama是transformer自然不改变形状,输出映射到词表做预测,词表32004长,第32004号词是

  4. 输出的序列进入MLP改变通道数,再抽取其中的进行SAM的提示编码,接着就是右边部分SAM的流程,参见《SAM2架构》,最终输出掩码和文本。

  • 可以看到LLAVA的图文输入是匹配的,有多少个目标要分割就要复制同样数量的图,而SAM的设计是只需要一张图,对应多个提示。并且分割领域,一张图通常有多个mask供训练,所以输入模型的batch通常是一张图采样的多个目标(mask),而真正的batch维度——也就是图像是通过循环实现输入,在图像编码等步骤是可以做到并行的,其它时候基本上不是并行的。

🚀token处理

图文token在输入llama之前有进行拼接等操作,这也是多模态大模型和普通大模型的区别所在,因为多了图像编码,原本的文本中会有<IMAGE>这种特殊词元代表图片,在输入LLM前需要把它们替换成对图像抽取的特征token。

  • input_ids的处理流程:
  1. 检测token:先在每个样本的input_ids中找出所有<image>的位置。如果根本没碰到,就只做普通的文本embedding并加入new_input_embeds

  2. 输入特征拼接:对于每个样本,遇到文本时正常编码,遇到图像的<image>会替换成图像编码。如果配置了<im_start><im_end>,还会把它们也做成embedding加入,如此进行以完成单个输入序列的构建;

  3. 长度对齐:把所有序列stack起来,如果不同样本的长度不同,就找到最大值max_len,对每行embedding前面或后面补0到相同长度。同时对attention_mask做相同的pad,使其形状变成[B,new_seq_len]。

  • labels(指的是标注数据的正确回答文本)的处理流程:
  1. 初始化:如果原始训练没有传 labels(即只做推理),就直接返回 None;否则,为每个样本构建一个新的 label 序列 cur_new_labels

  2. 每当你把一个 <IMAGE>换成 image_features 时,对应位置在 cur_new_labels 中就插入一个全值为 IGNORE_INDEX(比如 -100)的张量;这样训练时模型不会对这段 embedding 计算交叉熵 loss,因为那部分是图像,不是要预测的文本;

  3. 长度对齐:同上。

🚀训练技术

本节通过阅读train_ds.py脚本了解用到的技术,可以从中了解多模态大模型加载、deepspeed分布式训练、自定义目标层的LoRA等操作如何实现。

首先直接看main函数,先传参,加载分词器tokenizer,还会添加和标记图像编码开头和结尾的特殊词元,到时候合并图文编码时,图像编码会塞到这个两个词元的位置。

    args = parse_args(args)
    args.log_dir = os.path.join(args.log_base_dir, args.exp_name)
    if args.local_rank == 0:
        os.makedirs(args.log_dir, exist_ok=True)
        writer = SummaryWriter(args.log_dir)
    else:
        writer = None

    # Create model
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        args.version,
        cache_dir=None,
        model_max_length=args.model_max_length,
        padding_side="right",
        use_fast=False,
    )
    tokenizer.pad_token = tokenizer.unk_token
    num_added_tokens = tokenizer.add_tokens("[SEG]")
    args.seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]

    if args.use_mm_start_end:
        tokenizer.add_tokens(
            [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
        )

接下来是加载模型,半精度加载,vision_pretrained指的是视觉网络也就是SAM,vision_tower是LLAVA的视觉编码器。

    model_args = {
        "train_mask_decoder": args.train_mask_decoder,
        "out_dim": args.out_dim,
        "ce_loss_weight": args.ce_loss_weight,
        "dice_loss_weight": args.dice_loss_weight,
        "bce_loss_weight": args.bce_loss_weight,
        "seg_token_idx": args.seg_token_idx,
        "vision_pretrained": args.vision_pretrained,
        "vision_tower": args.vision_tower,
        "use_mm_start_end": args.use_mm_start_end,
    }
    torch_dtype = torch.float32
    if args.precision == "bf16":
        torch_dtype = torch.bfloat16
    elif args.precision == "fp16":
        torch_dtype = torch.half
    model = LISAForCausalLM.from_pretrained(
        args.version, torch_dtype=torch_dtype, low_cpu_mem_usage=True, **model_args
    )
    model.config.eos_token_id = tokenizer.eos_token_id
    model.config.bos_token_id = tokenizer.bos_token_id
    model.config.pad_token_id = tokenizer.pad_token_id

    model.enable_input_require_grads()
    model.gradient_checkpointing_enable()

    model.get_model().initialize_vision_modules(model.get_model().config)
    vision_tower = model.get_model().get_vision_tower()
    vision_tower.to(dtype=torch_dtype, device=args.local_rank)
    if not args.eval_only:
        model.get_model().initialize_lisa_modules(model.get_model().config)

    for p in vision_tower.parameters():
        p.requires_grad = False
    for p in model.get_model().mm_projector.parameters():
        p.requires_grad = False

    conversation_lib.default_conversation = conversation_lib.conv_templates[
        args.conv_type
    ]

接着加入LoRA权重find_linear_layers()是找出模型中特定层并返回层名的列表,这里选定的是线性层,并且排除了视觉相关模块,也就是不微调视觉编码的网络。args.lora_target_modules.split(",")选定的是”q_proj,v_proj”,也就是Q和V的矩阵加LoRA权重。

    lora_r = args.lora_r
    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_dropout = args.lora_dropout
        lora_target_modules = find_linear_layers(
            model, args.lora_target_modules.split(",")
        )
        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",
        )
        model = get_peft_model(model, lora_config)
        model.print_trainable_parameters()

把一些层的部分变得可训练,也就是这些层是全量微调的,torch.cuda.device_count()是有多少个GPU。

    model.resize_token_embeddings(len(tokenizer))

    # make text_hidden_fcs, mask_decoder, lm_head, embed_tokens trainable
    for n, p in model.named_parameters():
        if any(
            [
                x in n
                for x in ["lm_head", "embed_tokens", "mask_decoder", "text_hidden_fcs"]
            ]
        ):
            print("n: ", n, "p.shape: ", p.shape)
            p.requires_grad = True

    world_size = torch.cuda.device_count()
    args.distributed = world_size > 1

数据集加载,训练集会读取多个设定好的数据集,采样会随机抽,验证集则只支持refer_seg和reason_seg,数据集详情见下一节。

    train_dataset = HybridDataset(
        args.dataset_dir,
        tokenizer,
        args.vision_tower,
        samples_per_epoch=args.batch_size
        * args.grad_accumulation_steps
        * args.steps_per_epoch
        * world_size,
        precision=args.precision,
        image_size=args.image_size,
        num_classes_per_sample=args.num_classes_per_sample,
        exclude_val=args.exclude_val,
        dataset=args.dataset,
        sample_rate=[float(x) for x in args.sample_rates.split(",")],
        sem_seg_data=args.sem_seg_data,
        refer_seg_data=args.refer_seg_data,
        vqa_data=args.vqa_data,
        reason_seg_data=args.reason_seg_data,
        explanatory=args.explanatory,
    )

    if args.no_eval == False:
        val_dataset = ValDataset(
            args.dataset_dir,
            tokenizer,
            args.vision_tower,
            args.val_dataset,
            args.image_size,
        )
        print(
            f"Training with {len(train_dataset)} examples and validating with {len(val_dataset)} examples."
        )
    else:
        val_dataset = None
        print(f"Training with {len(train_dataset)} examples.")

deepspeed分布式训练器及其配置如下,大部分都是传参,优化器用的adamw,使用warmup学习率调度,也就是前100步学习率会慢慢上升至预设的学习率,后面就会线性下降到最小值。zero_optimization是deepspeed的zero优化,它把网络不同层的梯度存储和计算分配给不同GPU并行完成,优化了通信效率和显存压力。deepspeed.initialize()就会得到分布式的训练模型model_engine,它本质也是nn.Module,所以用起来和非分布式没什么区别。

Zero原理.png

  • Zero-0:不使用所有类型的分片,仅使用DeepSpeed作为DDP,速度最快(显存够时使用);

  • Zero-1:切分优化器状态,分片到每个数据并行的工作进程(每个GPU)下;有微小的速度提升;

  • Zero-2:切分优化器状态 + 梯度,分片到每个数据并行的工作进程(每个GPU)下;

  • Zero-3:切分优化器状态 + 梯度 + 模型参数,分片到每个数据并行的工作进程(每个GPU)下。

如上,Zero Stage从0-3训练速度越来越慢(不用offload),但训练所需显存递减。详情参考该文知乎

    ds_config = {
        "train_micro_batch_size_per_gpu": args.batch_size,
        "gradient_accumulation_steps": args.grad_accumulation_steps,
        "optimizer": {
            "type": "AdamW",
            "params": {
                "lr": args.lr,
                "weight_decay": 0.0,
                "betas": (args.beta1, args.beta2),
            },
        },
        "scheduler": {
            "type": "WarmupDecayLR",
            "params": {
                "total_num_steps": args.epochs * args.steps_per_epoch,
                "warmup_min_lr": 0,
                "warmup_max_lr": args.lr,
                "warmup_num_steps": 100,
                "warmup_type": "linear",
            },
        },
        "fp16": {
            "enabled": args.precision == "fp16",
        },
        "bf16": {
            "enabled": args.precision == "bf16",
        },
        "gradient_clipping": 1.0,
        "zero_optimization": {
            "stage": 2,
            "contiguous_gradients": True,
            "overlap_comm": True,
            "reduce_scatter": True,
            "reduce_bucket_size": 5e8,
            "allgather_bucket_size": 5e8,
        },
    }
    model_engine, optimizer, train_loader, scheduler = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        training_data=train_dataset,
        collate_fn=partial(
            collate_fn,
            tokenizer=tokenizer,
            conv_type=args.conv_type,
            use_mm_start_end=args.use_mm_start_end,
            local_rank=args.local_rank,
        ),
        config=ds_config,

然后是一些杂项,如果是恢复训练,就可以加载存好的检查点checkpoint。deepspeed会在优化器参数中存储模型的主参数,存储在global_step*/*optim_states.pt 文件中,数据类型为fp32。想要从checkpoint中恢复训练,不用动什么配置,直接重新启动程序即可。

然后把验证集设成分布式的,训练集不设置是因为本来就是随机抽样训练,每张卡各干各的就行了,而验证集就配合一起测完。

    # resume deepspeed checkpoint
    if args.auto_resume and len(args.resume) == 0:
        resume = os.path.join(args.log_dir, "ckpt_model")
        if os.path.exists(resume):
            args.resume = resume

    if args.resume:
        load_path, client_state = model_engine.load_checkpoint(args.resume)
        with open(os.path.join(args.resume, "latest"), "r") as f:
            ckpt_dir = f.readlines()[0].strip()
        args.start_epoch = (
            int(ckpt_dir.replace("global_step", "")) // args.steps_per_epoch
        )
        print(
            "resume training from {}, start from epoch {}".format(
                args.resume, args.start_epoch
            )
        )

    # validation dataset
    if val_dataset is not None:
        assert args.val_batch_size == 1
        val_sampler = torch.utils.data.distributed.DistributedSampler(
            val_dataset, shuffle=False, drop_last=False
        )
        val_loader = torch.utils.data.DataLoader(
            val_dataset,
            batch_size=args.val_batch_size,
            shuffle=False,
            num_workers=args.workers,
            pin_memory=False,
            sampler=val_sampler,
            collate_fn=partial(
                collate_fn,
                tokenizer=tokenizer,
                conv_type=args.conv_type,
                use_mm_start_end=args.use_mm_start_end,
                local_rank=args.local_rank,
            ),
        )

接着就进入训练和评估环节了:

    if args.eval_only:
        giou, ciou = validate(val_loader, model_engine, 0, writer, args)
        exit()

    for epoch in range(args.start_epoch, args.epochs):
        # train for one epoch
        train_iter = train(
            train_loader,
            model_engine,
            epoch,
            scheduler,
            writer,
            train_iter,
            args,
        )

        if args.no_eval == False:
            giou, ciou = validate(val_loader, model_engine, epoch, writer, args)
            is_best = giou > best_score
            best_score = max(giou, best_score)
            cur_ciou = ciou if is_best else cur_ciou

        if args.no_eval or is_best:
            save_dir = os.path.join(args.log_dir, "ckpt_model")
            if args.local_rank == 0:
                torch.save(
                    {"epoch": epoch},
                    os.path.join(
                        args.log_dir,
                        "meta_log_giou{:.3f}_ciou{:.3f}.pth".format(
                            best_score, cur_ciou
                        ),
                    ),
                )
                if os.path.exists(save_dir):
                    shutil.rmtree(save_dir)
            torch.distributed.barrier()
            model_engine.save_checkpoint(save_dir)

训练和评估的函数不多记录,和常规深度学习训练一样。可以看到deepspeed使得分布式训练和一般的训练使用起来没什么两样,这个框架所有算法需要关心的细节都浓缩到了配置项里面。

🔥数据集

🚀ade20k

ade20k数据集

  • 每张图都有若干个mask,每个mask的值对应一个分类,LISA在utils/sem_seg_dataset.py中的SemSegDataset为该数据集提供了处理方法,每张图被dataloader加载时,会随机选择该图包含的类别中的若干个,分别取其掩码,并生成对应的问答。也就是一图多个目标,每个目标对应一个类别、一个问答和一个掩码。这个数据集datasets类可以移植SAM2这边,但是需要去掉对话部分增加点作为提示。
class SemSegDataset(Dataset):
    pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    img_size = 1024
    ignore_label = 255

    def __init__(
        self,
        base_image_dir,
        image_size: int = 1024,
        num_classes_per_sample: int = 1,
        transform=None,
    ):
        self.transform = transform
        self.num_classes_per_sample = num_classes_per_sample

        self.base_image_dir = base_image_dir
        self.image_size = image_size

        classes, images, labels = self.init_ade20k(base_image_dir)
        self.data = (images, labels)
        self.classes = classes

    def init_ade20k(self, base_image_dir):
        with open("./ade20k_classes.json", "r") as f:
            ade20k_classes = json.load(f)
        ade20k_classes = np.array(ade20k_classes)
        image_ids = sorted(
            os.listdir(os.path.join(base_image_dir, "ade20k/images", "training"))
        )
        ade20k_image_ids = []
        for x in image_ids:
            if x.endswith(".jpg"):
                ade20k_image_ids.append(x[:-4])
        ade20k_images = []
        for image_id in ade20k_image_ids:  # self.descriptions:
            ade20k_images.append(
                os.path.join(
                    base_image_dir,
                    "ade20k",
                    "images",
                    "training",
                    "{}.jpg".format(image_id),
                )
            )
        ade20k_labels = [
            x.replace(".jpg", ".png").replace("images", "annotations")
            for x in ade20k_images
        ]
        print("ade20k: ", len(ade20k_images))
        return ade20k_classes, ade20k_images, ade20k_labels
    
    def __len__(self):
        return len(self.data[0])

    def preprocess(self, x: torch.Tensor) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def padding(self, x: torch.Tensor) -> torch.Tensor:
        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def __getitem__(self, idx):
        image, labels = self.data
        image_path = image[idx]
        label_path = labels[idx]
        
        label = Image.open(label_path)
        label = np.array(label)
        label[label == 0] = 255
        label -= 1
        label[label == 254] = 255

        img = cv2.imread(image_path)
        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # if self.transform:
        #     image = self.transform(image)
            
        unique_label = np.unique(label).tolist()            # 按类别分出所有mask
        if 255 in unique_label:
            unique_label.remove(255)
        if len(unique_label) == 0:
            return self.__getitem__(0)

        classes = [self.classes[class_id] for class_id in unique_label]     # 图中有的类名
        if len(classes) >= self.num_classes_per_sample:
            sampled_classes = np.random.choice(
                classes, size=self.num_classes_per_sample, replace=False
            ).tolist()
        else:
            sampled_classes = classes

        class_ids = []                  # 类别索引
        for sampled_cls in sampled_classes:
            class_id = self.classes.tolist().index(sampled_cls)
            class_ids.append(class_id)

        image = self.preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous())
        label = self.padding(torch.from_numpy(label).long())
        
        masks = []
        points = []
        points_label = []
        for class_id in class_ids:
            mask = label == class_id
            masks.append(mask)
            coords = torch.nonzero(mask, as_tuple=False)            # 坐标点
            if coords.numel() > 0:
                coords = coords.view(-1, 2)
                yx = coords[torch.randint(0, coords.shape[0], (1,)).item()]
                points.append([yx[1].item(), yx[0].item()])
                points_label.append(1)
            else:
                points.append([random.randint(0, 1024), random.randint(0, 1024)])
                points_label.append(0)
        
        masks = torch.stack(masks, dim=0)
        return (
            image,
            masks,
            np.array(points, dtype=np.float32),
            np.array(points_label, dtype=np.float32),
            # label,
            # sampled_classes,
            # image_path
        )

用SAM2微调一手(50个epoch)有提升,可以看到相较于2k样本左右的LabPicsV1数据集,这个20k的数据集指标就不可能那么高了:

规模 IoU GIoU CIoU Dice PA F1
SAM2(L) 0.514 0219 0.467 0.610 0.952 0.193
SAM2(L,ft)exp17 0.594 0.197 0.560 0.683 0.981 0.251

🚀COCO-stuff

这数据集结构应该和ade20k差不多,标注也是一样的。

COCO-stuff数据集

🚀refCOCO系列

refCOCO/refCOCO+/refCOCOg结构应该是一样的,接下来根据LISA读取数据集的相关函数进行debug我们可以看到其数据的构成。

  • 一个instances.json文件,记录了每个图像都信息images
{
  "license": 1,
  "file_name": "COCO_train2014_000000098304.jpg",
  "coco_url": "http://mscoco.org/images/98304",
  "height": 424,
  "width": 640,
  "date_captured": "2013-11-21 23:06:41",
  "flickr_url": "http://farm6.staticflickr.com/5062/5896644212_a326e96ea9_z.jpg",
  "id": 98304
}
  • 标注信息annotations,为所有目标(一个图会有多个目标)提供了边界框、mask等信息,掩码是记录多边形顶点:
{
  "segmentation": [[...]],
  "area": 197.29899999999986,
  "iscrowd": 0,
  "image_id": 98304,
  "bbox": [263.87, 216.88, 21.13, 15.17],
  "category_id": 18,
  "id": 3007
}
  • 80个分类categories
{
  "supercategory": "person",
  "id": 1,
  "name": "person"
}
  • 一个.p结尾的文件ref(unc).p,记录每个图的每个目标的描述文本,split表示属于训练集(train)、验证集(val):
{
  "sent_ids": [0, 1, 2],
  "file_name": "COCO_train2014_000000581857_16.jpg",
  "ann_id": 1719310,
  "ref_id": 0,
  "image_id": 581857,
  "split": "train",
  "sentences": [
    {
      "tokens": [...],
      "raw": "THE LADY WITH THE BLUE SHIRT",
      "sent_id": 0,
      "sent": "the lady with the blue shirt"
    },
    {
      "tokens": [...],
      "raw": "lady w back to us",
      "sent_id": 1,
      "sent": "lady with back to us"
    },
    {
      "tokens": [...],
      "raw": "blue shirt",
      "sent_id": 2,
      "sent": "blue shirt"
    }
  ],
  "category_id": 1
}

refclef也是和refcoco系列一样的,使用MSCOCO作为图像源,但是是比较早的,表达不够自然,偏人工合成风格。

🚀ReasonSeg

LISA做的数据集,就一千多张图,标注内容主要为推理文本和多边形标注的mask,也就是文本上不会直接提出目标的名字,而是提供推理的线索,如下图为例。其它其实和refCOCO差不多。

"the source of power for the ship"

字段名 类型 说明
text 字符串 问题文本,用于引导模型在图像中寻找相关区域进行回答。通常为完整英文句子。
is_sentence 布尔值 是否为完整句子;用于判断文本结构。
label 字符串 标注的名称,通常是目标类型或角色,例如 target 表示被问题指向的答案对象。
labels 字符串列表 多标签版本,一般包含和 label 一致的标签列表,用于兼容多标签识别。
shape_type 字符串 标注类型,例如 polygon 表示多边形框选区域。
points 二维浮点数组 多边形的顶点坐标(像素位置),每个点为 [x, y],顺序连接形成mask闭合区域。
image_name 字符串 图像文件名,标注所对应的图片,如 692198_ac99d18ac5_o.jpg
group_id 整数或 null 标注组 ID,用于把多个相关标注归为一组,若为 null 表示无分组。
group_ids 列表或 null 与 group_id 类似,用于支持多组 ID。
flags 字典 可选标记,常为空 {},也可包含如 "difficult": true 之类的信息。

🚀数据集确定

暂定如下数据集:

数据集结构

LISA上训练这些数据集是按照比例(sem_seg: refer_seg: vqa: reason_seg=9: 3: 3: 1)随机选取一个数据集抽样,每个图片会有多个目标及其对应分类,接着会从中随机抽取若干个类生成图文对(一个图对应多个目标,每个目标对应一个描述/问题)作为样本,对于没有描述文本的sem_seg数据集就会从预设的问题库中生成问答,有描述的如refer_seg就直接取用就行了。

🔥跑通

接下来这节记录使用数据集再次跑通训练、评估,会补充一些上一个文章没有的东西。

🚀微调

  • 训练时每个epoch会设定指定step数(也就是不用跑满每一个样本),LISA的梯度积累是每个step内传播grad_accumulation_steps次,所以单个epoch会用到step×grad_accumulation_steps×batch_size个图像,而不是step×batch_size个;总的mask个数则不超过step×grad_accumulation_steps×batch_size×num_classes_per_sample(真正并行的只有num_classes_per_sample和部分batch_size)。如果是多卡训练,这个batch size是每张卡的batch大小,所以样本数应当再乘以卡的数量。

  • 模型跑通训练如下,括号外是当前值(最新step),括号内是平均值。

微调输出

每个step都打印太冗长,我把它用tqdm改成动态进度条:

进度条

🚀评估

  • 上一篇文档用自定义的脚本和指标实现了对外部数据集的评估,而LISA的train_ds.py除了训练还提供了评估(支持官方用到的数据集),可以选择训练每一代后都接一个评估,也可以只评估或者只训练,由于该脚本的介绍已经讲过,而评估本身也只会更简单,所以就只补充一些没提到的内容;

  • 指标是gIoU和cIoU,前者是所有样本的IoU均值;后者则是积累所有样本的交集和并集面积后求IoU,它关注更大面积的目标分割效果。这是区别于GIoU和CIoU的(它们是IoU的变种)。

  • LISA提供的ValDataset如下(源码路径有点问题,修正了),评估只支持reason_seg和refer_seg的数据集,我们选择了后者,也就是refcoco、refcoco+、refcocog这仨。

  • __init__()就是把上面提到的图像、标注等信息读取出来。

class ValDataset(torch.utils.data.Dataset):
    pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    img_size = 1024
    ignore_label = 255

    def __init__(
        self,
        base_image_dir,
        tokenizer,
        vision_tower,
        val_dataset,
        image_size=1024,
    ):
        splits = val_dataset.split("|")
        self.base_image_dir = base_image_dir
        if len(splits) == 2:
            ds, split = splits
            images = glob.glob(
                os.path.join(self.base_image_dir, "reason_seg", ds, split, "*.jpg")
            )
            self.images = images
            self.data_type = "reason_seg"
        elif len(splits) == 3:
            ds, splitBy, split = splits
            refer_api = REFER(os.path.join(self.base_image_dir, "refer_seg"), ds, splitBy)
            ref_ids_val = refer_api.getRefIds(split=split)
            images_ids_val = refer_api.getImgIds(ref_ids=ref_ids_val)
            refs_val = refer_api.loadRefs(ref_ids=ref_ids_val)
            refer_seg_ds = {}
            refer_seg_ds["images"] = []
            loaded_images = refer_api.loadImgs(image_ids=images_ids_val)
            for item in loaded_images:
                item = item.copy()
                if ds == "refclef":
                    item["file_name"] = os.path.join(
                        self.base_image_dir, "refer_seg", "images/saiapr_tc-12", item["file_name"]
                    )
                elif ds in ["refcoco", "refcoco+", "refcocog", "grefcoco"]:
                    item["file_name"] = os.path.join(
                        self.base_image_dir, "refer_seg",
                        "images/mscoco/images/train2014",
                        item["file_name"],
                    )
                refer_seg_ds["images"].append(item)
            refer_seg_ds["annotations"] = refer_api.Anns  # anns_val

            img2refs = {}
            for ref in refs_val:
                image_id = ref["image_id"]
                img2refs[image_id] = img2refs.get(image_id, []) + [
                    ref,
                ]
            refer_seg_ds["img2refs"] = img2refs
            self.refer_seg_ds = refer_seg_ds
            self.data_type = "refer_seg"

        self.ds = ds
        self.image_size = image_size
        self.tokenizer = tokenizer
        self.transform = ResizeLongestSide(image_size)
        self.clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower)
  • __getitem__()内容和训练用的dataset差不多,也是会最终输出十个变量的样本,一张图对应多个目标,每个目标一个掩码和一个对话:
        return (
            image_path,      # 图像路径
            image,           # 图像
            image_clip,      # 变换后的图像
            conversations,   # 对话
            masks,           # 掩码,由多边形解码成能用的图像形式(一个图有多个掩码)
            labels,          # 整张图的标注(张量,每个像素属于哪一类/掩码)
            resize,          # 变换后的图像尺寸
            None,            # 问题(训练dataset有,这里不输出)
            None,            # 采样目标的分类名(训练dataset有,这里不输出)
            inference,       # 是否做推理任务
        )

🚀参数配置

微调和评估用到的部分参数(有些不用调或者很容易看懂用的就没写)如下,这只是单卡的所以仅供参考。

参数名 描述 参数值
version 预训练LLaVA权重路径,可以是LISA完整权重 ./weight/lisa
precision 精度(前向传播等操作) bf16
vision-tower LLaVA视觉主干权重路径 ./weight/clip-vit-large-patch14
dataset 训练的数据集类型,用 || 分隔多个,默认”sem_seg||refer_seg||vqa||reason_seg” sem_seg||refer_seg
sample_rates 每个类型数据集采样比重,数量和dataset一样,默认9,3,3,1 1,1
sem_seg_data sem_seg类型的数据集中挑选数据集,默认”ade20k||cocostuff||pascal_part||paco_lvis||mapillary” ade20k||cocostuff
refer_seg_data refer_seg类型的数据集中挑选数据集,默认”refclef||refcoco||refcoco+||refcocog” refcoco||refcoco+||refcocog
vqa_data vaq类型数据集,这里没用到就默认 llava_instruct_150k
reason_seg_data reason_seg类型数据集,没用到默认 ReasonSeg|train
val_dataset 验证(评估)数据集,默认ReasonSeg|val,如果是ref_seg应该是类似refcoco|unc|val,第一个是数据集名,第二个代表.p文件的编码方式(见.p文件名),第三个则是数据集切片(有val/test等) refcoco|unc|test
log_base_dir 记录每次训练用的文件夹,每次训练的tensorboard和权重都在里面 ./runs
steps_per_epoch 每一代多少步 1500
batch_size 批次大小(乘以steps_per_epoch就是一代采样数),显存不够只能1 1
grad_accumulation_steps 梯度积累 20
val_batch_size 验证集批次大小 1
num_classes_per_sample 每个样本(图像)取多少个类别(mask) 3
no_eval 不做验证为True False
eval_only 只做验证(评估)不训练为True True或False
vision_pretrained SAM预训练权重路径 ./weight/sam_vit_h_4b8939.pth

🚀复现

官方提供的7B模型评估结果(giou/ciou)如下,似乎和论文里面(ciou)的不一样,可能放出来的模型不是它们评估成绩最好的?不过原文分高的数据集在这分也确实是更高者。

又微调了一下refer_seg数据集,效果上去了,说明确实不是最优模型。

数据集→ refCOCO refCOCO+ refCOCOg
模型↓ val(1.5k) testA(750) testB(750) val(1.5k) testA(750) testB(750) val(1.3k) test(2.6k)
LISA-7B 67.3/66.6 70.3/70.1 63.0/62.3 52.8/53.4 59.6/59.5 45.1/45.9 60.3/60.2 61.2/61.3
LISA-7B(ft) 75.7/73.1 79.8/78.0 72.8/69.5 66.4/61.8 70.6/66.8 59.7/54.3 68.5/65.6 69.8/67.3

LISA论文的referseg数据集效果(ciou)

训练曲线

健康的损失曲线和评估曲线。

🚀多卡训练

目前单张4090D(24G)刚好跑满LISA-7B(BF16),而且只能batch size为1,跑18k个图片样本(实际上有更多的样本),1个epoch需要近2.8h,效率很低下。

要有足够的训练速度保证实验进度,保守需要总显存为48G(40G以上),所以硬件至少需要2张4090(或者等价的单张计算卡),如果有更多资源能应该能同时进行多个实验。

最终选择5张40G显存的A100

多卡训练输出

多卡训练如上图所示,多卡打印也会翻倍。运行时需要deepspeed --include localhost:3,4,5,6,7 train_ds.py指令指定需要使用的GPU的编号,这种情况下指定GPU并且不支持debug,所以debug应该单卡比较好。

多卡

速度问题解决:跑下来5张A100稳定时整体效率(相同时间输入的样本量)有单张4090D的3.5倍以上,显存是绰绰有余,但是训练依然要比较长的时间。

我发现训练时,后面的epoch花的时间比前面会少很多,刚开始1.5h每代,到后面甚至只要55分钟(此时效率大于5倍),可能是deepspeed框架会逐渐优化显存分配。最后发现++当8卡机剩余3个A100也在工作时,这5张卡速度也会上去++,AI的回答是(可能原因):

  • 部分机器上,GPU 之间通过 PCIe/NVLink 互联,当多个 GPU 同时有进程活跃时,驱动和硬件的 带宽调度机制会变得更积极,降低数据传输瓶颈。

  • 如果 GPU 0-2 也在运作,整个 PCIe 或 NVLink 互联模块会进入高负载状态 → 激发全局带宽/延迟优化 → 反而加速你 3-7 的训练过程。

所以理论上应该所有卡都在运行比较好。

于是我做了个实验,运行以下脚本来使用剩余3个GPU:

import torch
import time
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(torch.cuda.device_count())

# 模拟每张卡做一点小计算
tensors = []
for i in range(torch.cuda.device_count()):
    d = torch.device(f"cuda:{i}")
    # 创建一个小矩阵放在 GPU 上
    tensors.append(torch.randn(4, 4, device=d))

print("Starting light GPU occupation loop...")
try:
    while True:
        for i, tensor in enumerate(tensors):
            # 做一次小矩阵乘法 → 保证 GPU 有一点点算力占用
            result = tensor @ tensor.T
            _ = result.sum().item()  # 防止优化掉
        time.sleep(0.5)  # 每隔 0.5 秒运行一次
except KeyboardInterrupt:
    print("Stopped.")

发现不运行时训练程序会3.8s/step进行训练,而运行了脚本后能到达2.12s/step。所以这个脚本可以一直跑,以很低的负载让0-2GPU保持工作,这样即便有他人要用这三个GPU也不用担心有影响。

显存优化实验:我用一个小网络实验了一下,速度对比如下,所以再显存够的情况下stage0就行了,不需要做显存优化。

ZeRO Stage 平均步长时间 (s) 吞吐率 (samples/s)
0 0.0112 1427.72
1 0.0119 1341.57
2 0.0125 1279.72

batch设计:每张卡B(batch_size)为4时速度2.15s/step,而16时速度为7.65s/step左右,B翻了4倍但是费时仅为3倍多,因为GPU算大矩阵效率更高,所以B越大效率越高,但没必要拉满,要不然样本批次分布都差不多,结合GPT的建议设置为16-20为宜。

batch size选择(左图小,右图大

左边6样本每step,右边16*5样本每step,稳定性确实有很大区别。

参数推荐

参数名 参数值 参数名 参数值
zero 0 precision bf16
dataset sem_seg||refer_seg||reason_seg sample_rates 10,8,1
sem_seg_data ade20k||cocostuff refer_seg_data refclef||refcoco||refcoco+||refcocog
reason_seg_data ReasonSeg|train val_dataset refcoco|unc|testA
epochs 80 steps_per_epoch 500
batch_size 16 grad_accumulation_steps 1
lr 0.0003 ce_loss_weight 0.2
dice_loss_weight 1 bce_loss_weight 1.5

reason_seg数据集也就200+样本,所以应该比较低的采样比例,可能是因为这个数据集文本比较长,如果不巧抽到多个这个数据集的样本,可能会爆显存。

🔥问题解决

🚀SAM模型加载权重无效

加载模型时会出现这种情况:

/home/zigaa/anaconda3/envs/leosam2/lib/python3.10/site-packages/torch/nn/modules/module.py:2409: UserWarning: for mask_decoder.iou_prediction_head.layers.2.weight: copying from a non-meta parameter in the checkpoint to a meta parameter in the current model, which is a no-op. (Did you mean to pass assign=True to assign items in the state dictionary to their corresponding key in the module instead of copying them in place?)

而且不是上面这一行,而是所有SAM的模块都会报错,这是因为SAM模型构建时里有一些“meta 参数”(即占位符参数,没有在内存里真正分配张量),而你在加载检查点(checkpoint)时,试图把存储在检查点里的真实权重拷贝到这些meta参数上。因为meta参数本身并没有实际张量,所以“拷贝”操作变成了no‑op(无效操作),PyTorch就给打了Warning。

解决方法就是在build_sam()函数里面的sam.load_state_dict(state_dict, strict=False)添加assign=True,这会强制赋值权重文件的张量给meta参数的坑位。

这时候如果是评估不会有问题,但是训练时脚本执行到deepspeed.initialize()时会出现新的报错(截取最底层):

[rank0]:     work = group.broadcast([tensor], opts)
[rank0]: ValueError: Tensors must be contiguous

这是指那些参量虽然被强行赋值,但是依然不是真正分配在 CPU/GPU 上的张量,那些层很可能仍保持meta状态或被封装为非连续结构,就是空有数值而形式不对。所以应该把这些meta张量用一个新的contiguous张量替换。总的对build_sam()的改动如下:

def _build_sam(
    encoder_embed_dim,
    encoder_depth,
    encoder_num_heads,
    encoder_global_attn_indexes,
    checkpoint=None,
):
    prompt_embed_dim = 256
    image_size = 1024
    vit_patch_size = 16
    image_embedding_size = image_size // vit_patch_size
    sam = Sam(
        image_encoder=ImageEncoderViT(
            depth=encoder_depth,
            embed_dim=encoder_embed_dim,
            img_size=image_size,
            mlp_ratio=4,
            norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
            num_heads=encoder_num_heads,
            patch_size=vit_patch_size,
            qkv_bias=True,
            use_rel_pos=True,
            global_attn_indexes=encoder_global_attn_indexes,
            window_size=14,
            out_chans=prompt_embed_dim,
        ),
        prompt_encoder=PromptEncoder(
            embed_dim=prompt_embed_dim,
            image_embedding_size=(image_embedding_size, image_embedding_size),
            input_image_size=(image_size, image_size),
            mask_in_chans=16,
        ),
        mask_decoder=MaskDecoder(
            num_multimask_outputs=3,
            transformer=TwoWayTransformer(
                depth=2,
                embedding_dim=prompt_embed_dim,
                mlp_dim=2048,
                num_heads=8,
            ),
            transformer_dim=prompt_embed_dim,
            iou_head_depth=3,
            iou_head_hidden_dim=256,
        ),
        pixel_mean=[123.675, 116.28, 103.53],
        pixel_std=[58.395, 57.12, 57.375],
    )
    sam.eval()
    if checkpoint is not None:
        with open(checkpoint, "rb") as f:
            state_dict = torch.load(f)
        sam.load_state_dict(state_dict, strict=False, assign=True)

    from torch import nn
    device = torch.device("cuda")
    for name, param in sam.named_parameters():
        data = param.data
        # 1) 如果是 meta 张量,先创建一个同 shape/dtype 的真实张量
        if hasattr(data, "is_meta") and data.is_meta:
            real = torch.empty(
                data.shape,
                dtype=data.dtype,      # 保持原 dtype
                device=device          # 直接放到 CUDA
            )
            # 用 nn.Parameter 完全替换,确保整个 Parameter 对象都绑定到真实张量上
            new_p = nn.Parameter(real, requires_grad=param.requires_grad)
            # 把它插回到 model 里
            parent, attr = name.rsplit(".", 1)
            # 递归拿到父 module
            mod = sam
            for sub in parent.split("."):
                mod = getattr(mod, sub)
            setattr(mod, attr, new_p)
        # 2) 如果已经有真实张量,但是 non‑contiguous,强制 contiguous + clone + to(device)
        else:
            if not data.is_contiguous() or data.device != device:
                new_data = data.contiguous().clone().detach().to(device)
                param.data = new_data
    return sam

改动前微调模型输出的评估值通常会异常低且不收敛,这说明模型原本读取SAM部分的权重确实存在问题,改动完成后微调就能正常收敛。

🚀梯度积累问题

LISA梯度积累有问题,这种写法压根没有积累,依然是每次采样都进行更新,这样训练的效果相当不稳定,甚至它还没有清空梯度:

    for global_step in range(args.steps_per_epoch):
        for i in range(args.grad_accumulation_steps):
            try:
                input_dict = next(train_iter)
            except:
                train_iter = iter(train_loader)
                input_dict = next(train_iter)

            data_time.update(time.time() - end)
            input_dict = dict_to_cuda(input_dict)

            if args.precision == "fp16":
                input_dict["images"] = input_dict["images"].half()
                input_dict["images_clip"] = input_dict["images_clip"].half()
            elif args.precision == "bf16":
                input_dict["images"] = input_dict["images"].bfloat16()
                input_dict["images_clip"] = input_dict["images_clip"].bfloat16()
            else:
                input_dict["images"] = input_dict["images"].float()
                input_dict["images_clip"] = input_dict["images_clip"].float()

            output_dict = model(**input_dict)

            loss = output_dict["loss"]
            ce_loss = output_dict["ce_loss"]
            mask_bce_loss = output_dict["mask_bce_loss"]
            mask_dice_loss = output_dict["mask_dice_loss"]
            mask_loss = output_dict["mask_loss"]

            losses.update(loss.item(), input_dict["images"].size(0))
            ce_losses.update(ce_loss.item(), input_dict["images"].size(0))
            mask_bce_losses.update(mask_bce_loss.item(), input_dict["images"].size(0))
            mask_dice_losses.update(mask_dice_loss.item(), input_dict["images"].size(0))
            mask_losses.update(mask_loss.item(), input_dict["images"].size(0))
            model.backward(loss)
            model.step()

应该改成这样(就改后面部分),只反向传播得到梯度,到次数后在更新参数,清零梯度:

            losses.update(loss.item(), input_dict["images"].size(0))
            ce_losses.update(ce_loss.item(), input_dict["images"].size(0))
            mask_bce_losses.update(mask_bce_loss.item(), input_dict["images"].size(0))
            mask_dice_losses.update(mask_dice_loss.item(), input_dict["images"].size(0))
            mask_losses.update(mask_loss.item(), input_dict["images"].size(0))
        
            model.backward(loss)
            
        model.step()
        model.zero_grad()

🚀tensorboard记录问题

如下的这种tensorboard写法每个epoch都会重置step,从前面给到的源码可以看到global_step实际上并不是全局的(估计是编写的人没意识到),是单epoch的,所以应该用now_step = global_step + epoch * args.steps_per_epoch算出真正的全局step以替换global_step

            if args.local_rank == 0:
                progress.display(global_step + 1)
                writer.add_scalar("train/loss", losses.avg, global_step)
                writer.add_scalar("train/ce_loss", ce_losses.avg, global_step)
                writer.add_scalar(
                    "train/mask_bce_loss", mask_bce_losses.avg, global_step
                )
                writer.add_scalar(
                    "train/mask_dice_loss", mask_dice_losses.avg, global_step
                )
                writer.add_scalar("train/mask_loss", mask_losses.avg, global_step)
                writer.add_scalar(
                    "metrics/total_secs_per_batch", batch_time.avg, global_step
                )
                writer.add_scalar(
                    "metrics/data_secs_per_batch", data_time.avg, global_step
                )

🚀Zero1权重保存

Zero1方法训练时优化器参数不会分片,所以可以直接用如下代码代替zero_to_fp32.py的功能(这个脚本负责的是zero2和3)生成bin文件,然后再运行merge的脚本。

# zero1训练后获取完整权重(代替zero_to_fp32)
import torch

# 加载你的 checkpoint 文件路径
ckpt_path = "runs/lisa_feedback/ckpt_model/global_step22500/mp_rank_00_model_states.pt"

# 加载 checkpoint 文件
ckpt = torch.load(ckpt_path, map_location="cpu")

# 检查是否包含 'module' 这个关键 key(DeepSpeed Stage 1 保存方式)
if 'module' not in ckpt:
    raise KeyError("'module' key not found in checkpoint — 这不是一个合法的 DeepSpeed ZeRO Stage 1 模型参数文件。")

# 提取模型参数
state_dict = ckpt['module']

# 打印所有参数 key(检查是不是完整的,比如包含 q_proj.weight 等)
print("\n== 模型参数 key 列表(共 {} 个) ==".format(len(state_dict)))
for key in list(state_dict.keys())[:20]:  # 只打印前 20 个参数
    print(key)

# 也可以保存为临时的 pytorch_model.bin 以供 merge LoRA 使用
torch.save(state_dict, "pytorch_model.bin")
print("\n已保存模型参数到 pytorch_model.bin(可用于 merge_lora_weights_and_save_hf_model.py)")

🔥总结

本文在前文的基础上,使用LISA官方用的数据集复现了微调和评估,体验了多卡分布式训练,对数据集处理、基于deepspeed的分布式训练技术以及LISA本身的原理都有了很多深入的理解和认知,对多模态大模型的方方面面都有了初步的学习和理解。

  • Copyrights © 2023-2025 LegendLeo Chen
  • 访问人数: | 浏览次数:

请我喝杯咖啡吧~

支付宝
微信