import torch from torch.utils.data.distributed import DistributedSampler import torch.nn as nn # datasets related from lib.train.dataset import Lasot, Got10k, MSCOCOSeq, ImagenetVID, TrackingNet, Imagenet1k,VastTrack from lib.train.dataset import Lasot_lmdb, Got10k_lmdb, MSCOCOSeq_lmdb, ImagenetVID_lmdb, TrackingNet_lmdb from lib.train.dataset import VisEvent, LasHeR, DepthTrack from lib.train.dataset import Otb99_lang, Tnl2k, RefCOCOSeq,OTB_Lang from lib.train.data import sampler, opencv_loader, processing, LTRLoader import lib.train.data.transforms as tfm from lib.utils.misc import is_main_process from lib.train.dataset.refcoco_seq_JointNLT import RefCOCOSeq as RefCOCOSeq_jointnlt def update_settings(settings, cfg): settings.print_interval = cfg.TRAIN.PRINT_INTERVAL settings.search_area_factor = {'template': getattr(cfg.DATA.TEMPLATE, "FACTOR", None), 'search': getattr(cfg.DATA.SEARCH, "FACTOR", None)} settings.output_sz = {'template': getattr(cfg.DATA.TEMPLATE, "SIZE", 128), 'search': getattr(cfg.DATA.SEARCH, "SIZE", 256)} settings.center_jitter_factor = {'template': getattr(cfg.DATA.TEMPLATE, "CENTER_JITTER", None), 'search':getattr(cfg.DATA.SEARCH, "CENTER_JITTER", None)} settings.scale_jitter_factor = {'template': getattr(cfg.DATA.TEMPLATE, "SCALE_JITTER", None), 'search': getattr(cfg.DATA.SEARCH, "SCALE_JITTER", None)} settings.grad_clip_norm = cfg.TRAIN.GRAD_CLIP_NORM settings.print_stats = None settings.batchsize = cfg.TRAIN.BATCH_SIZE settings.scheduler_type = cfg.TRAIN.SCHEDULER.TYPE settings.multi_modal_vision = getattr(cfg.DATA, "MULTI_MODAL_VISION", False) settings.multi_modal_language = getattr(cfg.DATA, "MULTI_MODAL_LANGUAGE", False) train_type = getattr(cfg.TRAIN, "TYPE", None) if train_type == "peft": settings.fix_norm = True else: settings.fix_norm = False def names2datasets(name_list: list, settings, image_loader): assert isinstance(name_list, list) datasets = [] for name in name_list: assert name in ["LASOT", "GOT10K_vottrain", "GOT10K_votval", "GOT10K_train_full", "COCO17", "VID", "TRACKINGNET", "IMAGENET1K", "DepthTrack_train", "DepthTrack_val", "LasHeR_all", "LasHeR_train","LasHeR_val", "VisEvent", "REFCOCOG", "TNL2K_train", "OTB99_train","OTB_Lang", "VastTrack",'RefCOCO14'] if name == "LASOT": if settings.use_lmdb: print("Building lasot dataset from lmdb") datasets.append(Lasot_lmdb(settings.env.lasot_lmdb_dir, split='train', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language)) else: datasets.append(Lasot(settings.env.lasot_dir, split='train', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language)) if name == "GOT10K_vottrain": if settings.use_lmdb: print("Building got10k from lmdb") datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='vottrain', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) else: datasets.append(Got10k(settings.env.got10k_dir, split='vottrain', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "GOT10K_train_full": if settings.use_lmdb: print("Building got10k_train_full from lmdb") datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='train_full', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) else: datasets.append(Got10k(settings.env.got10k_dir, split='train_full', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "GOT10K_votval": if settings.use_lmdb: print("Building got10k from lmdb") datasets.append(Got10k_lmdb(settings.env.got10k_lmdb_dir, split='votval', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) else: datasets.append(Got10k(settings.env.got10k_dir, split='votval', image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "COCO17": if settings.use_lmdb: print("Building COCO2017 from lmdb") datasets.append(MSCOCOSeq_lmdb(settings.env.coco_lmdb_dir, version="2017", image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) else: datasets.append(MSCOCOSeq(settings.env.coco_dir, version="2017", image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "VID": if settings.use_lmdb: print("Building VID from lmdb") datasets.append(ImagenetVID_lmdb(settings.env.imagenet_lmdb_dir, image_loader=image_loader)) else: datasets.append(ImagenetVID(settings.env.imagenet_dir, image_loader=image_loader)) if name == "TRACKINGNET": if settings.use_lmdb: print("Building TrackingNet from lmdb") datasets.append(TrackingNet_lmdb(settings.env.trackingnet_lmdb_dir, image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) else: # raise ValueError("NOW WE CAN ONLY USE TRACKINGNET FROM LMDB") datasets.append(TrackingNet(settings.env.trackingnet_dir, image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "IMAGENET1K": datasets.append(Imagenet1k(settings.env.imagenet1k_dir, image_loader=image_loader)) if name == "DepthTrack_train": datasets.append(DepthTrack(settings.env.depthtrack_dir, dtype='rgbcolormap', split='train', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "DepthTrack_val": datasets.append(DepthTrack(settings.env.depthtrack_dir, dtype='rgbcolormap', split='val', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "LasHeR_all": datasets.append(LasHeR(settings.env.lasher_dir, dtype='rgbrgb', split='all', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "LasHeR_train": datasets.append(LasHeR(settings.env.lasher_dir, dtype='rgbrgb', split='train', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "LasHeR_val": datasets.append(LasHeR(settings.env.lasher_dir, dtype='rgbrgb', split='val', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "VisEvent": datasets.append(VisEvent(settings.env.visevent_dir, dtype='rgbrgb', split='train', multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "REFCOCOG": datasets.append(RefCOCOSeq(settings.env.refcoco_dir, split="train", image_loader=image_loader, name="refcocog", splitBy="google", multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) if name == "RefCOCO14": # datasets.append(RefCOCOSeq(settings.env.ref_coco_dir, refcoco_type="refcoco-unc", version="2014", image_loader=image_loader)) datasets.append(RefCOCOSeq_jointnlt(settings.env.ref_coco_dir, split="train", image_loader=image_loader, name="refcocog", splitBy='google')) if name == "TNL2K_train": datasets.append(Tnl2k(settings.env.tnl2k_dir, split=None, image_loader=image_loader, multi_modal_vision=settings.multi_modal_vision, multi_modal_language=settings.multi_modal_language )) elif name == "OTB99_train": # datasets.append(Otb99_lang(settings.env.otb99_dir, split='train', image_loader=image_loader, # multi_modal_vision=settings.multi_modal_vision, # multi_modal_language=settings.multi_modal_language # )) datasets.append(OTB_Lang(settings.env.otb99_dir, split='train', image_loader=image_loader)) if name == "VastTrack": datasets.append(VastTrack(settings.env.vasttrack_dir, split='train', image_loader=image_loader)) return datasets def build_dataloaders(cfg, settings): settings.num_template = getattr(cfg.DATA.TEMPLATE, "NUMBER", 1) settings.num_search = getattr(cfg.DATA.SEARCH, "NUMBER", 1) # Data transform transform_joint = tfm.Transform(tfm.ToGrayscale(probability=0.05), tfm.RandomHorizontalFlip(probability=0.5)) transform_train = tfm.Transform(tfm.ToTensorAndJitter(0.2), tfm.RandomHorizontalFlip_Norm(probability=0.5), tfm.Normalize(mean=cfg.DATA.MEAN, std=cfg.DATA.STD)) # The tracking pairs processing module output_sz = settings.output_sz search_area_factor = settings.search_area_factor data_processing_train = processing.SeqTrackProcessing(search_area_factor=search_area_factor, output_sz=output_sz, center_jitter_factor=settings.center_jitter_factor, scale_jitter_factor=settings.scale_jitter_factor, mode='sequence', transform=transform_train, joint_transform=transform_joint, multi_modal_language=settings.multi_modal_language, settings=settings) # Train sampler and loader sampler_mode = getattr(cfg.DATA, "SAMPLER_MODE", "causal") dataset_train = sampler.TrackingSampler(datasets=names2datasets(cfg.DATA.TRAIN.DATASETS_NAME, settings, opencv_loader), p_datasets=cfg.DATA.TRAIN.DATASETS_RATIO, samples_per_epoch=cfg.DATA.TRAIN.SAMPLE_PER_EPOCH, max_gap=cfg.DATA.MAX_SAMPLE_INTERVAL, num_search_frames=settings.num_search, num_template_frames=settings.num_template, processing=data_processing_train, frame_sample_mode=sampler_mode ) train_sampler = DistributedSampler(dataset_train) if settings.local_rank != -1 else None shuffle = False if settings.local_rank != -1 else True loader_train = LTRLoader('train', dataset_train, training=True, batch_size=cfg.TRAIN.BATCH_SIZE, shuffle=shuffle, num_workers=cfg.TRAIN.NUM_WORKER, drop_last=True, stack_dim=1, sampler=train_sampler) return loader_train def get_optimizer_scheduler(net, cfg): train_type = getattr(cfg.TRAIN, "TYPE", None) if train_type == "target_state": trainable_keywords = getattr(cfg.TRAIN, "TARGET_STATE_TRAINABLE", [ "target_state_encoder.projector", "target_state_encoder.film_ln", "target_state_encoder.film", "target_state_encoder.film_gate", "lora_", "box_head", "confidence_pred", ]) # Keywords for params that start from scratch and need higher LR qwen_lr_keywords = [ "target_state_encoder.projector", "target_state_encoder.film_ln", "target_state_encoder.film", "target_state_encoder.film_gate", "lora_", ] for n, p in net.named_parameters(): p.requires_grad = any(key in n for key in trainable_keywords) target_encoder = getattr(net, "target_state_encoder", None) if target_encoder is not None: target_encoder.configure_token_embedding_training( getattr(cfg.MODEL.TARGET_STATE, "TRAIN_TOKEN_EMBEDDING", False) ) qwen_embedding_param = None if target_encoder is not None: try: qwen_embedding_param = target_encoder.qwen.get_input_embeddings().weight except AttributeError: qwen_embedding_param = None qwen_lr_mult = getattr(cfg.TRAIN, "QWEN_LR_MULTIPLIER", 1.0) base_lr = cfg.TRAIN.LR tracker_params = [] # box_head, confidence_pred → baseline LR qwen_params = [] # projector, film, lora → higher LR embedding_params = [] # TARGET_STATE embedding row → baseline LR, wd=0 for n, p in net.named_parameters(): if not p.requires_grad: continue if qwen_embedding_param is not None and p is qwen_embedding_param: embedding_params.append(p) elif any(key in n for key in qwen_lr_keywords): qwen_params.append(p) else: tracker_params.append(p) param_dicts = [] if tracker_params: param_dicts.append({"params": tracker_params, "lr": base_lr}) if qwen_params: param_dicts.append({"params": qwen_params, "lr": base_lr * qwen_lr_mult}) if embedding_params: param_dicts.append({"params": embedding_params, "lr": base_lr, "weight_decay": 0.0}) if is_main_process(): print("Learnable parameters are shown below.") print(f" (base_lr={base_lr}, qwen_lr={base_lr * qwen_lr_mult})") for n, p in net.named_parameters(): if p.requires_grad: print(f" {n}") elif train_type == "peft": param_dicts = [ {"params": [p for n, p in net.named_parameters() if "prompt" in n or "interface" in n and p.requires_grad]}, ] for n, p in net.named_parameters(): if ("prompt" not in n) and ("interface" not in n): p.requires_grad = False if is_main_process(): print("Learnable parameters are shown below.") for n, p in net.named_parameters(): if p.requires_grad: print(n) else: for n, p in net.named_parameters(): if 'text_encoder' in n and p.requires_grad: p.requires_grad = False print("Freeze: ", n) # for n, p in net.named_parameters(): # if 'backbone' in n and p.requires_grad: # p.requires_grad = False # print("Freeze: ", n) param_dicts = [ {"params": [p for n, p in net.named_parameters() if "backbone" not in n and p.requires_grad]}, { "params": [p for n, p in net.named_parameters() if "backbone" in n and p.requires_grad], "lr": cfg.TRAIN.LR * cfg.TRAIN.ENCODER_MULTIPLIER, }, ] train_n_list = [] if is_main_process(): print("Learnable parameters are shown below.") for n, p in net.named_parameters(): if p.requires_grad: train_n_list.append(n) print(n) if cfg.TRAIN.OPTIMIZER == "ADAMW": optimizer = torch.optim.AdamW(param_dicts, lr=cfg.TRAIN.LR, weight_decay=cfg.TRAIN.WEIGHT_DECAY) else: raise ValueError("Unsupported Optimizer") # ---- LR scheduler ---- scheduler_type = cfg.TRAIN.SCHEDULER.TYPE if scheduler_type == 'step': lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, cfg.TRAIN.LR_DROP_EPOCH ) elif scheduler_type == "Mstep": lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=cfg.TRAIN.SCHEDULER.MILESTONES, gamma=cfg.TRAIN.SCHEDULER.GAMMA, ) elif scheduler_type == "cosine": # Cosine annealing with optional linear warmup. import math as _math warmup_epochs = getattr(cfg.TRAIN, "WARMUP_EPOCHS", 0) total_epochs = cfg.TRAIN.EPOCH T_max = max(1, total_epochs - warmup_epochs) if warmup_epochs > 0: def _lr_lambda(epoch): if epoch < warmup_epochs: return float(epoch + 1) / float(max(1, warmup_epochs)) progress = float(epoch - warmup_epochs) / float(T_max) return 0.5 * (1.0 + _math.cos(_math.pi * progress)) else: def _lr_lambda(epoch): progress = float(epoch) / float(max(1, total_epochs)) return 0.5 * (1.0 + _math.cos(_math.pi * progress)) lr_scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=_lr_lambda ) else: raise ValueError(f"Unsupported scheduler type: {scheduler_type}") return optimizer, lr_scheduler