import argparse import os import shutil import sys import time from functools import partial import deepspeed import torch import tqdm import transformers from peft import LoraConfig, get_peft_model from torch.utils.tensorboard import SummaryWriter from VisualSearch.model.VSM import VSMForCausalLM from VisualSearch.model.llava import conversation as conversation_lib from VisualSearch.utils.dataset import HybridDataset, ValDataset, collate_fn from VisualSearch.utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, AverageMeter, ProgressMeter, Summary, dict_to_cuda, intersectionAndUnionGPU) def parse_args(args): parser = argparse.ArgumentParser(description="VisualSearch Model Training") parser.add_argument("--local_rank", default=0, type=int, help="node rank") parser.add_argument( "--version", default="LLaVA-7B-v1.1" ) parser.add_argument( "--precision", default="bf16", type=str, choices=["fp32", "bf16", "fp16"], help="precision for training", ) parser.add_argument("--model_max_length", default=512, type=int) parser.add_argument("--lora_r", default=8, type=int) parser.add_argument( "--vision-tower", default="openai/clip-vit-large-patch14", type=str ) parser.add_argument("--load_in_8bit", action="store_true", default=False) parser.add_argument("--load_in_4bit", action="store_true", default=False) parser.add_argument( "--dataset", default="general_segdet||refer_seg||mixed_grounding||vqa", type=str ) parser.add_argument("--sample_rates", default="15,4,4,15", type=str) parser.add_argument( "--general_segdet_data", default="objects365||cocostuff||paco_lvis", type=str, ) parser.add_argument("--general_segdet_sample_rates", default="2,1,1", type=str) parser.add_argument( "--refer_seg_data", default="refclef||refcoco||refcoco+||refcocog", type=str ) parser.add_argument("--vqa_data", default="possible_locations_conv_86k||llava_instruct_80k", type=str) parser.add_argument("--vqa_sample_rates", default="2,1", type=str) parser.add_argument("--val_dataset", default="refcoco|unc|val", type=str) parser.add_argument("--dataset_dir", default="data", type=str) parser.add_argument("--log_base_dir", default="./runs", type=str) parser.add_argument("--exp_name", default="vsm", type=str) parser.add_argument("--epochs", default=40, type=int) parser.add_argument("--steps_per_epoch", default=2500, type=int) parser.add_argument( "--batch_size", default=4, type=int, help="batch size per device per step" ) parser.add_argument( "--grad_accumulation_steps", default=2, type=int, ) parser.add_argument("--val_batch_size", default=1, type=int) parser.add_argument("--workers", default=2, type=int) parser.add_argument("--lr", default=0.0001, type=float) parser.add_argument("--ce_loss_weight", default=1.0, type=float) parser.add_argument("--dice_loss_weight", default=0.5, type=float) parser.add_argument("--bce_loss_weight", default=2.0, type=float) parser.add_argument("--det_loss_weight", default=0.1, type=float) parser.add_argument("--lora_alpha", default=16, type=int) parser.add_argument("--lora_dropout", default=0.05, type=float) parser.add_argument("--lora_target_modules", default="q_proj,v_proj", type=str) parser.add_argument("--explanatory", default=0.1, type=float) parser.add_argument("--beta1", default=0.9, type=float) parser.add_argument("--beta2", default=0.95, type=float) parser.add_argument("--num_classes_per_sample", default=3, type=int) parser.add_argument("--exclude_val", action="store_true", default=False) parser.add_argument("--no_eval", action="store_true", default=False) parser.add_argument("--out_dim", default=512, type=int) parser.add_argument("--weight", type=str) parser.add_argument("--resume", default="", type=str) parser.add_argument("--print_freq", default=1, type=int) parser.add_argument("--start_epoch", default=0, type=int) parser.add_argument("--gradient_checkpointing", action="store_true", default=True) parser.add_argument("--train_mask_decoder", action="store_true", default=True) parser.add_argument("--use_mm_start_end", action="store_true", default=True) parser.add_argument("--auto_resume", action="store_true", default=False) parser.add_argument( "--conv_type", default="llava_v1", type=str, choices=["llava_v1", "llava_llama_2"], ) return parser.parse_args(args) def box_cxcywh_to_xyxy(x): x_c, y_c, w, h = x.unbind(1) b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] return torch.stack(b, dim=1) def iou(bbox1, bbox2): x1 = max(bbox1[0], bbox2[0]) y1 = max(bbox1[1], bbox2[1]) x2 = min(bbox1[2], bbox2[2]) y2 = min(bbox1[3], bbox2[3]) w1 = bbox1[2] - bbox1[0] h1 = bbox1[3] - bbox1[1] w2 = bbox2[2] - bbox2[0] h2 = bbox2[3] - bbox2[1] inter_area = max(0, x2 - x1) * max(0, y2 - y1) return inter_area/(w1*h1+w2*h2-inter_area) def main(args): 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("[LOC]") args.loc_token_idx = tokenizer("[LOC]", 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 ) 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, "det_loss_weight" : args.det_loss_weight, "loc_token_idx": args.loc_token_idx, "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 = VSMForCausalLM.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) 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 = True conversation_lib.default_conversation = conversation_lib.conv_templates[ args.conv_type ] 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 [ "owlvit", "visual_projection", "prompt_encoder", "mask_decoder", "vision_tower", "mm_projector", "text_hidden_fcs_seg", "text_hidden_fcs_det", ] ] ) 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() 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", "visual_projection", "prompt_encoder", "mask_decoder", "text_hidden_fcs_seg", "text_hidden_fcs_det", "owlvit.class_head", "owlvit.layer_norm"] ] ): # print("n: ", n, "p.shape: ", p.shape) p.requires_grad = True world_size = torch.cuda.device_count() print('world_size', world_size) args.distributed = world_size > 1 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, 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(",")], general_segdet_data=args.general_segdet_data, general_segdet_sample_rate=[float(x) for x in args.general_segdet_sample_rates.split(",")], refer_seg_data=args.refer_seg_data, vqa_data=args.vqa_data, vqa_sample_rate=[float(x) for x in args.vqa_sample_rates.split(",")], ) if args.no_eval == False: val_dataset = ValDataset( args.dataset_dir, tokenizer, args.vision_tower, args.val_dataset, ) print( f"Training with {len(train_dataset)} examples and validating with {len(val_dataset)} examples." ) 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, ) # 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, 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, ), ) train_iter = iter(train_loader) best_score, cur_ciou, cur_giou = 0.0, 0.0, 0.0 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, det_acc = validate(val_loader, model_engine, epoch, writer, args) is_best = det_acc > best_score best_score = max(det_acc, best_score) cur_giou = giou if is_best else cur_giou 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_detacc{:.3f}_giou{:.3f}_ciou{:.3f}.pth".format( best_score, cur_giou, cur_ciou ), ), ) if os.path.exists(save_dir): shutil.rmtree(save_dir) torch.distributed.barrier() model_engine.save_checkpoint(save_dir) def train( train_loader, model, epoch, scheduler, writer, train_iter, args, ): """Main training loop.""" batch_time = AverageMeter("Time", ":6.3f") data_time = AverageMeter("Data", ":6.3f") losses = AverageMeter("Loss", ":.4f") ce_losses = AverageMeter("CeLoss", ":.4f") mask_bce_losses = AverageMeter("MaskBCELoss", ":.4f") mask_dice_losses = AverageMeter("MaskDICELoss", ":.4f") mask_losses = AverageMeter("MaskLoss", ":.4f") detection_losses = AverageMeter("DetectionLoss", ":.4f") detection_ce_losses = AverageMeter("DetectionCELoss", ":.4f") detection_bbox_losses = AverageMeter("DetectionBBOXLoss", ":.4f") detection_giou_losses = AverageMeter("DetectionGIOULoss", ":.4f") progress = ProgressMeter( args.steps_per_epoch, [ batch_time, losses, ce_losses, mask_losses, mask_bce_losses, mask_dice_losses, detection_losses, detection_ce_losses, detection_bbox_losses, detection_giou_losses ], prefix="Epoch: [{}]".format(epoch), ) # switch to train mode model.train() end = time.time() 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"] detection_loss = output_dict['detection_loss'] detection_loss_ce = output_dict['detection_loss_ce'] detection_loss_bbox = output_dict['detection_loss_bbox'] detection_loss_giou = output_dict['detection_loss_giou'] losses.update(loss.item(), 1) ce_losses.update(ce_loss.item(), 1) mask_bce_losses.update(mask_bce_loss.item(), 1) mask_dice_losses.update(mask_dice_loss.item(), 1) mask_losses.update(mask_loss.item(), 1) detection_losses.update(detection_loss.item(), 1) detection_ce_losses.update(detection_loss_ce.item(), 1) detection_bbox_losses.update(detection_loss_bbox.item(), 1) detection_giou_losses.update(detection_loss_giou.item(), 1) model.backward(loss) model.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() if global_step % args.print_freq == 0: if args.distributed: batch_time.all_reduce() data_time.all_reduce() losses.all_reduce() ce_losses.all_reduce() mask_bce_losses.all_reduce() mask_dice_losses.all_reduce() mask_losses.all_reduce() detection_losses.all_reduce() detection_ce_losses.all_reduce() detection_bbox_losses.all_reduce() detection_giou_losses.all_reduce() if args.local_rank == 0: progress.display(global_step + 1) writer.add_scalar("train/loss", losses.avg, global_step+args.steps_per_epoch*epoch) writer.add_scalar("train/ce_loss", ce_losses.avg, global_step+args.steps_per_epoch*epoch) writer.add_scalar( "train/mask_bce_loss", mask_bce_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "train/mask_dice_loss", mask_dice_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar("train/mask_loss", mask_losses.avg, global_step+args.steps_per_epoch*epoch) writer.add_scalar( "train/detection_loss", detection_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "train/detection_ce_loss", detection_ce_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "train/detection_bbox_loss", detection_bbox_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "train/detection_giou_loss", detection_giou_losses.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "metrics/total_secs_per_batch", batch_time.avg, global_step+args.steps_per_epoch*epoch ) writer.add_scalar( "metrics/data_secs_per_batch", data_time.avg, global_step+args.steps_per_epoch*epoch ) batch_time.reset() data_time.reset() losses.reset() ce_losses.reset() mask_bce_losses.reset() mask_dice_losses.reset() mask_losses.reset() detection_losses.reset() detection_ce_losses.reset() detection_bbox_losses.reset() detection_giou_losses.reset() if global_step != 0: curr_lr = scheduler.get_last_lr() if args.local_rank == 0: writer.add_scalar("train/lr", curr_lr[0], global_step+args.steps_per_epoch*epoch) return train_iter def validate(val_loader, model_engine, epoch, writer, args): intersection_meter = AverageMeter("Intersec", ":6.3f", Summary.SUM) union_meter = AverageMeter("Union", ":6.3f", Summary.SUM) acc_iou_meter = AverageMeter("gIoU", ":6.3f", Summary.SUM) det_acc_meter = AverageMeter("DetAcc", ":6.3f", Summary.SUM) model_engine.eval() for input_dict in tqdm.tqdm(val_loader): torch.cuda.empty_cache() 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() with torch.no_grad(): output_dict = model_engine(**input_dict) pred_masks = output_dict["pred_masks"] masks_list = output_dict["gt_masks"][0].int() output_list = (pred_masks[0] > 0).int() assert len(pred_masks) == 1 pred_logits = output_dict['pred_logits'] pred_boxes = output_dict['pred_boxes'] gt_bboxes = output_dict['gt_bboxes'] for pred_logits_i, pred_boxes_i, gt_bboxes_i in zip(pred_logits, pred_boxes, gt_bboxes): top_index = pred_logits_i.view(-1).argmax() pred_bbox = pred_boxes_i[top_index].view(1, 4) gt_bbox = gt_bboxes_i.view(1,4) iou_i = iou(box_cxcywh_to_xyxy(pred_bbox).view(4), box_cxcywh_to_xyxy(gt_bbox).view(4)) det_acc = 1.0 if iou_i > 0.5 else 0.0 det_acc_meter.update(det_acc, 1) intersection, union, acc_iou = 0.0, 0.0, 0.0 for mask_i, output_i in zip(masks_list, output_list): intersection_i, union_i, _ = intersectionAndUnionGPU( output_i.contiguous().clone(), mask_i.contiguous(), 2, ignore_index=255 ) intersection += intersection_i union += union_i acc_iou += intersection_i / (union_i + 1e-5) acc_iou[union_i == 0] += 1.0 # no-object target intersection, union = intersection.cpu().numpy(), union.cpu().numpy() acc_iou = acc_iou.cpu().numpy() / masks_list.shape[0] intersection_meter.update(intersection), union_meter.update( union ), acc_iou_meter.update(acc_iou, n=masks_list.shape[0]) intersection_meter.all_reduce() union_meter.all_reduce() acc_iou_meter.all_reduce() det_acc_meter.all_reduce() iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) ciou = iou_class[1] giou = acc_iou_meter.avg[1] det_acc = det_acc_meter.avg if args.local_rank == 0: writer.add_scalar("val/giou", giou, epoch) writer.add_scalar("val/ciou", ciou, epoch) writer.add_scalar("val/det_acc", det_acc, epoch) print("giou: {:.4f}, ciou: {:.4f}, det_acc: {:.4f}".format(giou, ciou, det_acc)) return giou, ciou, det_acc if __name__ == "__main__": main(sys.argv[1:])