init

lib/train/run_training.py

+import os
+import sys
+import argparse
+import importlib
+import cv2 as cv
+import torch.backends.cudnn
+import torch.distributed as dist
+import torch
+import random
+import numpy as np
+torch.backends.cudnn.benchmark = False
+
+import _init_paths
+import lib.train.admin.settings as ws_settings
+
+
+def init_seeds(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    torch.set_num_threads(4)
+    cv.setNumThreads(1)
+    cv.ocl.setUseOpenCL(False)
+
+
+def run_training(script_name, config_name, cudnn_benchmark=True, local_rank=-1, save_dir=None, base_seed=None,
+                 use_lmdb=False, script_name_prv=None, config_name_prv=None, use_wandb=False,
+                 distill=None, script_teacher=None, config_teacher=None):
+    """Run the train script.
+    args:
+        script_name: Name of emperiment in the "experiments/" folder.
+        config_name: Name of the yaml file in the "experiments/<script_name>".
+        cudnn_benchmark: Use cudnn benchmark or not (default is True).
+    """
+    if save_dir is None:
+        print("save_dir dir is not given. Use the default dir instead.")
+    # This is needed to avoid strange crashes related to opencv
+    torch.set_num_threads(4)
+    cv.setNumThreads(4)
+
+    torch.backends.cudnn.benchmark = cudnn_benchmark
+
+    print('script_name: {}.py  config_name: {}.yaml'.format(script_name, config_name))
+
+    '''2021.1.5 set seed for different process'''
+    if base_seed is not None:
+        if local_rank != -1:
+            init_seeds(base_seed + local_rank)
+        else:
+            init_seeds(base_seed)
+
+    settings = ws_settings.Settings()
+    settings.script_name = script_name
+    settings.config_name = config_name
+    settings.project_path = 'train/{}/{}'.format(script_name, config_name)
+    if script_name_prv is not None and config_name_prv is not None:
+        settings.project_path_prv = 'train/{}/{}'.format(script_name_prv, config_name_prv)
+    settings.local_rank = local_rank
+    settings.save_dir = os.path.abspath(save_dir)
+    settings.use_lmdb = use_lmdb
+    prj_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
+    settings.cfg_file = os.path.join(prj_dir, 'experiments/%s/%s.yaml' % (script_name, config_name))
+    settings.use_wandb = use_wandb
+    if distill:
+        settings.distill = distill
+        settings.script_teacher = script_teacher
+        settings.config_teacher = config_teacher
+        if script_teacher is not None and config_teacher is not None:
+            settings.project_path_teacher = 'train/{}/{}'.format(script_teacher, config_teacher)
+        settings.cfg_file_teacher = os.path.join(prj_dir, 'experiments/%s/%s.yaml' % (script_teacher, config_teacher))
+        expr_module = importlib.import_module('lib.train.train_script_distill')
+    else:
+        expr_module = importlib.import_module('lib.train.train_script')
+    expr_func = getattr(expr_module, 'run')
+
+    expr_func(settings)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Run a train scripts in train_settings.')
+    parser.add_argument('--script', type=str, required=True, help='Name of the train script.')
+    parser.add_argument('--config', type=str, required=True, help="Name of the config file.")
+    parser.add_argument('--cudnn_benchmark', type=bool, default=False, help='Set cudnn benchmark on (1) or off (0) (default is on).')
+    parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
+    parser.add_argument('--save_dir', type=str, help='the directory to save checkpoints and logs')
+    parser.add_argument('--seed', type=int, default=42, help='seed for random numbers')
+    parser.add_argument('--use_lmdb', type=int, choices=[0, 1], default=0)  # whether datasets are in lmdb format
+    parser.add_argument('--script_prv', type=str, default=None, help='Name of the train script of previous model.')
+    parser.add_argument('--config_prv', type=str, default=None, help="Name of the config file of previous model.")
+    parser.add_argument('--use_wandb', type=int, choices=[0, 1], default=0)  # whether to use wandb
+    # for knowledge distillation
+    parser.add_argument('--distill', type=int, choices=[0, 1], default=0)  # whether to use knowledge distillation
+    parser.add_argument('--script_teacher', type=str, help='teacher script name')
+    parser.add_argument('--config_teacher', type=str, help='teacher yaml configure file name')
+
+    args = parser.parse_args()
+    if args.local_rank != -1:
+        dist.init_process_group(backend='nccl')
+        torch.cuda.set_device(args.local_rank)
+    else:
+        torch.cuda.set_device(0)
+    run_training(args.script, args.config, cudnn_benchmark=args.cudnn_benchmark,
+                 local_rank=args.local_rank, save_dir=args.save_dir, base_seed=args.seed,
+                 use_lmdb=args.use_lmdb, script_name_prv=args.script_prv, config_name_prv=args.config_prv,
+                 use_wandb=args.use_wandb,
+                 distill=args.distill, script_teacher=args.script_teacher, config_teacher=args.config_teacher)
+
+
+if __name__ == '__main__':
+    main()

lib/train/train_script.py

+import os
+# loss function related
+from lib.utils.box_ops import giou_loss
+from torch.nn.functional import l1_loss
+from torch.nn import BCEWithLogitsLoss
+# train pipeline related
+from lib.train.trainers import LTRTrainer, LTRSeqTrainer
+from lib.train.dataset import Lasot, Got10k, MSCOCOSeq, ImagenetVID, TrackingNet
+from lib.train.dataset import Lasot_lmdb, Got10k_lmdb, MSCOCOSeq_lmdb, ImagenetVID_lmdb, TrackingNet_lmdb
+from lib.train.data import sampler, opencv_loader, processing, LTRLoader, sequence_sampler
+# distributed training related
+from torch.nn.parallel import DistributedDataParallel as DDP
+# some more advanced functions
+from .base_functions import *
+# network related
+from lib.models.artrack import build_artrack
+from lib.models.artrack_seq import build_artrack_seq
+# forward propagation related
+from lib.train.actors import ARTrackActor, ARTrackSeqActor
+# for import modules
+import importlib
+
+from ..utils.focal_loss import FocalLoss
+
+def names2datasets(name_list: list, settings, image_loader):
+    assert isinstance(name_list, list)
+    datasets = []
+    #settings.use_lmdb = True
+    for name in name_list:
+        assert name in ["LASOT", "GOT10K_vottrain", "GOT10K_votval", "GOT10K_train_full", "GOT10K_official_val",
+                        "COCO17", "VID", "TRACKINGNET"]
+        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))
+            else:
+                datasets.append(Lasot(settings.env.lasot_dir, split='train', image_loader=image_loader))
+        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))
+            else:
+                datasets.append(Got10k(settings.env.got10k_dir, split='vottrain', image_loader=image_loader))
+        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))
+            else:
+                datasets.append(Got10k(settings.env.got10k_dir, split='train_full', image_loader=image_loader))
+        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))
+            else:
+                datasets.append(Got10k(settings.env.got10k_dir, split='votval', image_loader=image_loader))
+        if name == "GOT10K_official_val":
+            if settings.use_lmdb:
+                raise ValueError("Not implement")
+            else:
+                datasets.append(Got10k(settings.env.got10k_val_dir, split=None, image_loader=image_loader))
+        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))
+            else:
+                datasets.append(MSCOCOSeq(settings.env.coco_dir, version="2017", image_loader=image_loader))
+        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))
+            else:
+                # raise ValueError("NOW WE CAN ONLY USE TRACKINGNET FROM LMDB")
+                datasets.append(TrackingNet(settings.env.trackingnet_dir, image_loader=image_loader))
+    return datasets
+
+def slt_collate(batch):
+    ret = {}
+    for k in batch[0].keys():
+        here_list = []
+        for ex in batch:
+            here_list.append(ex[k])
+        ret[k] = here_list
+    return ret
+
+class SLTLoader(torch.utils.data.dataloader.DataLoader):
+    """
+    Data loader. Combines a dataset and a sampler, and provides
+    single- or multi-process iterators over the dataset.
+    """
+
+    __initialized = False
+
+    def __init__(self, name, dataset, training=True, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,
+                 num_workers=0, epoch_interval=1, collate_fn=None, stack_dim=0, pin_memory=False, drop_last=False,
+                 timeout=0, worker_init_fn=None):
+
+        if collate_fn is None:
+            collate_fn = slt_collate
+
+        super(SLTLoader, self).__init__(dataset, batch_size, shuffle, sampler, batch_sampler,
+                 num_workers, collate_fn, pin_memory, drop_last,
+                 timeout, worker_init_fn)
+
+        self.name = name
+        self.training = training
+        self.epoch_interval = epoch_interval
+        self.stack_dim = stack_dim
+
+def run(settings):
+    settings.description = 'Training script for STARK-S, STARK-ST stage1, and STARK-ST stage2'
+
+    # update the default configs with config file
+    if not os.path.exists(settings.cfg_file):
+        raise ValueError("%s doesn't exist." % settings.cfg_file)
+    config_module = importlib.import_module("lib.config.%s.config" % settings.script_name)
+    cfg = config_module.cfg
+    config_module.update_config_from_file(settings.cfg_file)
+    if settings.local_rank in [-1, 0]:
+        print("New configuration is shown below.")
+        for key in cfg.keys():
+            print("%s configuration:" % key, cfg[key])
+            print('\n')
+
+    # update settings based on cfg
+    update_settings(settings, cfg)
+
+    # Record the training log
+    log_dir = os.path.join(settings.save_dir, 'logs')
+    if settings.local_rank in [-1, 0]:
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+    settings.log_file = os.path.join(log_dir, "%s-%s.log" % (settings.script_name, settings.config_name))
+
+    # Build dataloaders
+
+    if "RepVGG" in cfg.MODEL.BACKBONE.TYPE or "swin" in cfg.MODEL.BACKBONE.TYPE or "LightTrack" in cfg.MODEL.BACKBONE.TYPE:
+        cfg.ckpt_dir = settings.save_dir
+    bins = cfg.MODEL.BINS
+    search_size = cfg.DATA.SEARCH.SIZE
+    # Create network
+    if settings.script_name == "artrack":
+        net = build_artrack(cfg)
+        loader_train, loader_val = build_dataloaders(cfg, settings)
+    elif settings.script_name == "artrack_seq":
+        net = build_artrack_seq(cfg)
+        dataset_train = sequence_sampler.SequenceSampler(
+            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_GAP, max_interval=cfg.DATA.MAX_INTERVAL,
+            num_search_frames=cfg.DATA.SEARCH.NUMBER, num_template_frames=1,
+            frame_sample_mode='random_interval',
+            prob=cfg.DATA.INTERVAL_PROB)
+        loader_train = SLTLoader('train', dataset_train, training=True, batch_size=cfg.TRAIN.BATCH_SIZE,
+                                 num_workers=cfg.TRAIN.NUM_WORKER,
+                                 shuffle=False, drop_last=True)
+    else:
+        raise ValueError("illegal script name")
+
+    # wrap networks to distributed one
+    net.cuda()
+    if settings.local_rank != -1:
+        # net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)  # add syncBN converter
+        net = DDP(net, device_ids=[settings.local_rank], find_unused_parameters=True)
+        settings.device = torch.device("cuda:%d" % settings.local_rank)
+    else:
+        settings.device = torch.device("cuda:0")
+    settings.deep_sup = getattr(cfg.TRAIN, "DEEP_SUPERVISION", False)
+    settings.distill = getattr(cfg.TRAIN, "DISTILL", False)
+    settings.distill_loss_type = getattr(cfg.TRAIN, "DISTILL_LOSS_TYPE", "KL")
+    # Loss functions and Actors
+    if settings.script_name == "artrack":
+        focal_loss = FocalLoss()
+        objective = {'giou': giou_loss, 'l1': l1_loss, 'focal': focal_loss}
+        loss_weight = {'giou': cfg.TRAIN.GIOU_WEIGHT, 'l1': cfg.TRAIN.L1_WEIGHT, 'focal': 2.}
+        actor = ARTrackActor(net=net, objective=objective, loss_weight=loss_weight, settings=settings, cfg=cfg, bins=bins, search_size=search_size)
+    elif settings.script_name == "artrack_seq":
+        focal_loss = FocalLoss()
+        objective = {'giou': giou_loss, 'l1': l1_loss, 'focal': focal_loss}
+        loss_weight = {'giou': cfg.TRAIN.GIOU_WEIGHT, 'l1': cfg.TRAIN.L1_WEIGHT, 'focal': 2.}
+        actor = ARTrackSeqActor(net=net, objective=objective, loss_weight=loss_weight, settings=settings, cfg=cfg, bins=bins, search_size=search_size)
+    else:
+        raise ValueError("illegal script name")
+
+    # if cfg.TRAIN.DEEP_SUPERVISION:
+    #     raise ValueError("Deep supervision is not supported now.")
+
+    # Optimizer, parameters, and learning rates
+    optimizer, lr_scheduler = get_optimizer_scheduler(net, cfg)
+    use_amp = getattr(cfg.TRAIN, "AMP", False)
+    if settings.script_name == "artrack":
+        trainer = LTRTrainer(actor, [loader_train, loader_val], optimizer, settings, lr_scheduler, use_amp=use_amp)
+    elif settings.script_name == "artrack_seq":
+        trainer = LTRSeqTrainer(actor, [loader_train], optimizer, settings, lr_scheduler, use_amp=use_amp)
+
+    # train process
+    trainer.train(cfg.TRAIN.EPOCH, load_latest=True, fail_safe=True)

lib/train/train_script_distill.py

+import os
+# loss function related
+from lib.utils.box_ops import giou_loss
+from torch.nn.functional import l1_loss
+from torch.nn import BCEWithLogitsLoss
+# train pipeline related
+from lib.train.trainers import LTRTrainer
+# distributed training related
+from torch.nn.parallel import DistributedDataParallel as DDP
+# some more advanced functions
+from .base_functions import *
+# network related
+from lib.models.stark import build_starks, build_starkst
+from lib.models.stark import build_stark_lightning_x_trt
+# forward propagation related
+from lib.train.actors import STARKLightningXtrtdistillActor
+# for import modules
+import importlib
+
+
+def build_network(script_name, cfg):
+    # Create network
+    if script_name == "stark_s":
+        net = build_starks(cfg)
+    elif script_name == "stark_st1" or script_name == "stark_st2":
+        net = build_starkst(cfg)
+    elif script_name == "stark_lightning_X_trt":
+        net = build_stark_lightning_x_trt(cfg, phase="train")
+    else:
+        raise ValueError("illegal script name")
+    return net
+
+
+def run(settings):
+    settings.description = 'Training script for STARK-S, STARK-ST stage1, and STARK-ST stage2'
+
+    # update the default configs with config file
+    if not os.path.exists(settings.cfg_file):
+        raise ValueError("%s doesn't exist." % settings.cfg_file)
+    config_module = importlib.import_module("lib.config.%s.config" % settings.script_name)
+    cfg = config_module.cfg
+    config_module.update_config_from_file(settings.cfg_file)
+    if settings.local_rank in [-1, 0]:
+        print("New configuration is shown below.")
+        for key in cfg.keys():
+            print("%s configuration:" % key, cfg[key])
+            print('\n')
+
+    # update the default teacher configs with teacher config file
+    if not os.path.exists(settings.cfg_file_teacher):
+        raise ValueError("%s doesn't exist." % settings.cfg_file_teacher)
+    config_module_teacher = importlib.import_module("lib.config.%s.config" % settings.script_teacher)
+    cfg_teacher = config_module_teacher.cfg
+    config_module_teacher.update_config_from_file(settings.cfg_file_teacher)
+    if settings.local_rank in [-1, 0]:
+        print("New teacher configuration is shown below.")
+        for key in cfg_teacher.keys():
+            print("%s configuration:" % key, cfg_teacher[key])
+            print('\n')
+
+    # update settings based on cfg
+    update_settings(settings, cfg)
+
+    # Record the training log
+    log_dir = os.path.join(settings.save_dir, 'logs')
+    if settings.local_rank in [-1, 0]:
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+    settings.log_file = os.path.join(log_dir, "%s-%s.log" % (settings.script_name, settings.config_name))
+
+    # Build dataloaders
+    loader_train, loader_val = build_dataloaders(cfg, settings)
+
+    if "RepVGG" in cfg.MODEL.BACKBONE.TYPE or "swin" in cfg.MODEL.BACKBONE.TYPE:
+        cfg.ckpt_dir = settings.save_dir
+    """turn on the distillation mode"""
+    cfg.TRAIN.DISTILL = True
+    cfg_teacher.TRAIN.DISTILL = True
+    net = build_network(settings.script_name, cfg)
+    net_teacher = build_network(settings.script_teacher, cfg_teacher)
+
+    # wrap networks to distributed one
+    net.cuda()
+    net_teacher.cuda()
+    net_teacher.eval()
+
+    if settings.local_rank != -1:
+        net = DDP(net, device_ids=[settings.local_rank], find_unused_parameters=True)
+        net_teacher = DDP(net_teacher, device_ids=[settings.local_rank], find_unused_parameters=True)
+        settings.device = torch.device("cuda:%d" % settings.local_rank)
+    else:
+        settings.device = torch.device("cuda:0")
+    # settings.deep_sup = getattr(cfg.TRAIN, "DEEP_SUPERVISION", False)
+    # settings.distill = getattr(cfg.TRAIN, "DISTILL", False)
+    settings.distill_loss_type = getattr(cfg.TRAIN, "DISTILL_LOSS_TYPE", "L1")
+    # Loss functions and Actors
+    if settings.script_name == "stark_lightning_X_trt":
+        objective = {'giou': giou_loss, 'l1': l1_loss}
+        loss_weight = {'giou': cfg.TRAIN.GIOU_WEIGHT, 'l1': cfg.TRAIN.L1_WEIGHT}
+        actor = STARKLightningXtrtdistillActor(net=net, objective=objective, loss_weight=loss_weight, settings=settings,
+                                               net_teacher=net_teacher)
+    else:
+        raise ValueError("illegal script name")
+
+    # Optimizer, parameters, and learning rates
+    optimizer, lr_scheduler = get_optimizer_scheduler(net, cfg)
+    use_amp = getattr(cfg.TRAIN, "AMP", False)
+    trainer = LTRTrainer(actor, [loader_train, loader_val], optimizer, settings, lr_scheduler, use_amp=use_amp)
+
+    # train process
+    trainer.train(cfg.TRAIN.EPOCH, load_latest=True, fail_safe=True, distill=True)

lib/train/trainers/__init__.py

+from .base_trainer import BaseTrainer
+from .ltr_trainer import LTRTrainer
+from .ltr_seq_trainer import LTRSeqTrainer

lib/train/trainers/base_trainer.py

+import os
+import glob
+import torch
+import traceback
+from lib.train.admin import multigpu
+from torch.utils.data.distributed import DistributedSampler
+
+
+class BaseTrainer:
+    """Base trainer class. Contains functions for training and saving/loading checkpoints.
+    Trainer classes should inherit from this one and overload the train_epoch function."""
+
+    def __init__(self, actor, loaders, optimizer, settings, lr_scheduler=None):
+        """
+        args:
+            actor - The actor for training the network
+            loaders - list of dataset loaders, e.g. [train_loader, val_loader]. In each epoch, the trainer runs one
+                        epoch for each loader.
+            optimizer - The optimizer used for training, e.g. Adam
+            settings - Training settings
+            lr_scheduler - Learning rate scheduler
+        """
+        self.actor = actor
+        self.optimizer = optimizer
+        self.lr_scheduler = lr_scheduler
+        self.loaders = loaders
+
+        self.update_settings(settings)
+
+        self.epoch = 0
+        self.stats = {}
+
+        self.device = getattr(settings, 'device', None)
+        if self.device is None:
+            self.device = torch.device("cuda:0" if torch.cuda.is_available() and settings.use_gpu else "cpu")
+
+        self.actor.to(self.device)
+        self.settings = settings
+
+    def update_settings(self, settings=None):
+        """Updates the trainer settings. Must be called to update internal settings."""
+        if settings is not None:
+            self.settings = settings
+
+        if self.settings.env.workspace_dir is not None:
+            self.settings.env.workspace_dir = os.path.expanduser(self.settings.env.workspace_dir)
+            '''2021.1.4 New function: specify checkpoint dir'''
+            if self.settings.save_dir is None:
+                self._checkpoint_dir = os.path.join(self.settings.env.workspace_dir, 'checkpoints')
+            else:
+                self._checkpoint_dir = os.path.join(self.settings.save_dir, 'checkpoints')
+            print("checkpoints will be saved to %s" % self._checkpoint_dir)
+
+            if self.settings.local_rank in [-1, 0]:
+                if not os.path.exists(self._checkpoint_dir):
+                    print("Training with multiple GPUs. checkpoints directory doesn't exist. "
+                          "Create checkpoints directory")
+                    os.makedirs(self._checkpoint_dir)
+        else:
+            self._checkpoint_dir = None
+
+    def train(self, max_epochs, load_latest=False, fail_safe=True, load_previous_ckpt=False, distill=False):
+        """Do training for the given number of epochs.
+        args:
+            max_epochs - Max number of training epochs,
+            load_latest - Bool indicating whether to resume from latest epoch.
+            fail_safe - Bool indicating whether the training to automatically restart in case of any crashes.
+        """
+
+        epoch = -1
+        num_tries = 1
+        for i in range(num_tries):
+            try:
+                if load_latest:
+                    self.load_checkpoint()
+                if load_previous_ckpt:
+                    directory = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path_prv)
+                    self.load_state_dict(directory)
+                if distill:
+                    directory_teacher = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path_teacher)
+                    self.load_state_dict(directory_teacher, distill=True)
+                for epoch in range(self.epoch+1, max_epochs+1):
+                    self.epoch = epoch
+
+                    self.train_epoch()
+
+                    if self.lr_scheduler is not None:
+                        if self.settings.scheduler_type != 'cosine':
+                            self.lr_scheduler.step()
+                        else:
+                            self.lr_scheduler.step(epoch - 1)
+                    # only save the last 10 checkpoints
+                    save_every_epoch = getattr(self.settings, "save_every_epoch", False)
+                    save_epochs = []
+                    if epoch > (max_epochs - 1) or save_every_epoch or epoch % 5 == 0 or epoch in save_epochs or epoch > (max_epochs - 5):
+                    # if epoch > (max_epochs - 10) or save_every_epoch or epoch % 100 == 0:
+                        if self._checkpoint_dir:
+                            if self.settings.local_rank in [-1, 0]:
+                                self.save_checkpoint()
+            except:
+                print('Training crashed at epoch {}'.format(epoch))
+                if fail_safe:
+                    self.epoch -= 1
+                    load_latest = True
+                    print('Traceback for the error!')
+                    print(traceback.format_exc())
+                    print('Restarting training from last epoch ...')
+                else:
+                    raise
+
+        print('Finished training!')
+
+    def train_epoch(self):
+        raise NotImplementedError
+
+    def save_checkpoint(self):
+        """Saves a checkpoint of the network and other variables."""
+
+        net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net
+
+        actor_type = type(self.actor).__name__
+        net_type = type(net).__name__
+        state = {
+            'epoch': self.epoch,
+            'actor_type': actor_type,
+            'net_type': net_type,
+            'net': net.state_dict(),
+            'net_info': getattr(net, 'info', None),
+            'constructor': getattr(net, 'constructor', None),
+            'optimizer': self.optimizer.state_dict(),
+            'stats': self.stats,
+            'settings': self.settings
+        }
+
+        directory = '{}/{}'.format(self._checkpoint_dir, self.settings.project_path)
+        print(directory)
+        if not os.path.exists(directory):
+            print("directory doesn't exist. creating...")
+            os.makedirs(directory)
+
+        # First save as a tmp file
+        tmp_file_path = '{}/{}_ep{:04d}.tmp'.format(directory, net_type, self.epoch)
+        torch.save(state, tmp_file_path)
+
+        file_path = '{}/{}_ep{:04d}.pth.tar'.format(directory, net_type, self.epoch)
+
+        # Now rename to actual checkpoint. os.rename seems to be atomic if files are on same filesystem. Not 100% sure
+        os.rename(tmp_file_path, file_path)
+
+    def load_checkpoint(self, checkpoint = None, fields = None, ignore_fields = None, load_constructor = False):
+        """Loads a network checkpoint file.
+
+        Can be called in three different ways:
+            load_checkpoint():
+                Loads the latest epoch from the workspace. Use this to continue training.
+            load_checkpoint(epoch_num):
+                Loads the network at the given epoch number (int).
+            load_checkpoint(path_to_checkpoint):
+                Loads the file from the given absolute path (str).
+        """
+
+        net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net
+
+        actor_type = type(self.actor).__name__
+        net_type = type(net).__name__
+
+        if checkpoint is None:
+            # Load most recent checkpoint
+            checkpoint_list = sorted(glob.glob('{}/{}/{}_ep*.pth.tar'.format(self._checkpoint_dir,
+                                                                             self.settings.project_path, net_type)))
+            if checkpoint_list:
+                checkpoint_path = checkpoint_list[-1]
+            else:
+                print('No matching checkpoint file found')
+                return
+        elif isinstance(checkpoint, int):
+            # Checkpoint is the epoch number
+            checkpoint_path = '{}/{}/{}_ep{:04d}.pth.tar'.format(self._checkpoint_dir, self.settings.project_path,
+                                                                 net_type, checkpoint)
+        elif isinstance(checkpoint, str):
+            # checkpoint is the path
+            if os.path.isdir(checkpoint):
+                checkpoint_list = sorted(glob.glob('{}/*_ep*.pth.tar'.format(checkpoint)))
+                if checkpoint_list:
+                    checkpoint_path = checkpoint_list[-1]
+                else:
+                    raise Exception('No checkpoint found')
+            else:
+                checkpoint_path = os.path.expanduser(checkpoint)
+        else:
+            raise TypeError
+
+        # Load network
+        checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+
+        assert net_type == checkpoint_dict['net_type'], 'Network is not of correct type.'
+
+        if fields is None:
+            fields = checkpoint_dict.keys()
+        if ignore_fields is None:
+            ignore_fields = ['settings']
+
+            # Never load the scheduler. It exists in older checkpoints.
+        ignore_fields.extend(['lr_scheduler', 'constructor', 'net_type', 'actor_type', 'net_info'])
+
+        # Load all fields
+        for key in fields:
+            if key in ignore_fields:
+                continue
+            if key == 'net':
+                net.load_state_dict(checkpoint_dict[key])
+            elif key == 'optimizer':
+                self.optimizer.load_state_dict(checkpoint_dict[key])
+            else:
+                setattr(self, key, checkpoint_dict[key])
+
+        # Set the net info
+        if load_constructor and 'constructor' in checkpoint_dict and checkpoint_dict['constructor'] is not None:
+            net.constructor = checkpoint_dict['constructor']
+        if 'net_info' in checkpoint_dict and checkpoint_dict['net_info'] is not None:
+            net.info = checkpoint_dict['net_info']
+
+        # Update the epoch in lr scheduler
+        if 'epoch' in fields:
+            self.lr_scheduler.last_epoch = self.epoch
+        # 2021.1.10 Update the epoch in data_samplers
+            for loader in self.loaders:
+                if isinstance(loader.sampler, DistributedSampler):
+                    loader.sampler.set_epoch(self.epoch)
+        return True
+
+    def load_state_dict(self, checkpoint=None, distill=False):
+        """Loads a network checkpoint file.
+
+        Can be called in three different ways:
+            load_checkpoint():
+                Loads the latest epoch from the workspace. Use this to continue training.
+            load_checkpoint(epoch_num):
+                Loads the network at the given epoch number (int).
+            load_checkpoint(path_to_checkpoint):
+                Loads the file from the given absolute path (str).
+        """
+        if distill:
+            net = self.actor.net_teacher.module if multigpu.is_multi_gpu(self.actor.net_teacher) \
+                else self.actor.net_teacher
+        else:
+            net = self.actor.net.module if multigpu.is_multi_gpu(self.actor.net) else self.actor.net
+
+        net_type = type(net).__name__
+
+        if isinstance(checkpoint, str):
+            # checkpoint is the path
+            if os.path.isdir(checkpoint):
+                checkpoint_list = sorted(glob.glob('{}/*_ep*.pth.tar'.format(checkpoint)))
+                if checkpoint_list:
+                    checkpoint_path = checkpoint_list[-1]
+                else:
+                    raise Exception('No checkpoint found')
+            else:
+                checkpoint_path = os.path.expanduser(checkpoint)
+        else:
+            raise TypeError
+
+        # Load network
+        print("Loading pretrained model from ", checkpoint_path)
+        checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
+
+        assert net_type == checkpoint_dict['net_type'], 'Network is not of correct type.'
+
+        missing_k, unexpected_k = net.load_state_dict(checkpoint_dict["net"], strict=False)
+        print("previous checkpoint is loaded.")
+        print("missing keys: ", missing_k)
+        print("unexpected keys:", unexpected_k)
+
+        return True

lib/train/trainers/ltr_seq_trainer.py

+import os
+import datetime
+from collections import OrderedDict
+from torch.nn.utils import clip_grad_norm_
+# from lib.train.data.wandb_logger import WandbWriter
+from lib.train.trainers import BaseTrainer
+from lib.train.admin import AverageMeter, StatValue
+from memory_profiler import profile
+# from lib.train.admin import TensorboardWriter
+import torch
+import time
+import numpy as np
+from torch.utils.data.distributed import DistributedSampler
+from torch.cuda.amp import autocast
+from torch.cuda.amp import GradScaler
+
+from lib.utils.misc import get_world_size
+
+
+class LTRSeqTrainer(BaseTrainer):
+    def __init__(self, actor, loaders, optimizer, settings, lr_scheduler=None, use_amp=False):
+        """
+        args:
+            actor - The actor for training the network
+            loaders - list of dataset loaders, e.g. [train_loader, val_loader]. In each epoch, the trainer runs one
+                        epoch for each loader.
+            optimizer - The optimizer used for training, e.g. Adam
+            settings - Training settings
+            lr_scheduler - Learning rate scheduler
+        """
+        super().__init__(actor, loaders, optimizer, settings, lr_scheduler)
+
+        self._set_default_settings()
+
+        # Initialize statistics variables
+        self.stats = OrderedDict({loader.name: None for loader in self.loaders})
+
+        # Initialize tensorboard and wandb
+        # self.wandb_writer = None
+        # if settings.local_rank in [-1, 0]:
+        #    tensorboard_writer_dir = os.path.join(self.settings.env.tensorboard_dir, self.settings.project_path)
+        #    if not os.path.exists(tensorboard_writer_dir):
+        #        os.makedirs(tensorboard_writer_dir)
+        #    self.tensorboard_writer = TensorboardWriter(tensorboard_writer_dir, [l.name for l in loaders])
+
+        #    if settings.use_wandb:
+        #        world_size = get_world_size()
+        #        cur_train_samples = self.loaders[0].dataset.samples_per_epoch * max(0, self.epoch - 1)
+        #        interval = (world_size * settings.batchsize)  # * interval
+        #        self.wandb_writer = WandbWriter(settings.project_path[6:], {}, tensorboard_writer_dir, cur_train_samples, interval)
+
+        self.move_data_to_gpu = getattr(settings, 'move_data_to_gpu', True)
+        print("move_data", self.move_data_to_gpu)
+        self.settings = settings
+        self.use_amp = use_amp
+        if use_amp:
+            self.scaler = GradScaler()
+
+    def _set_default_settings(self):
+        # Dict of all default values
+        default = {'print_interval': 10,
+                   'print_stats': None,
+                   'description': ''}
+
+        for param, default_value in default.items():
+            if getattr(self.settings, param, None) is None:
+                setattr(self.settings, param, default_value)
+
+        self.miou_list = []
+
+    def cycle_dataset(self, loader):
+        """Do a cycle of training or validation."""
+        torch.autograd.set_detect_anomaly(True)
+        self.actor.train(loader.training)
+        torch.set_grad_enabled(loader.training)
+
+        self._init_timing()
+
+        for i, data in enumerate(loader, 1):
+            self.actor.eval()
+            self.data_read_done_time = time.time()
+            with torch.no_grad():
+                explore_result = self.actor.explore(data)
+            if explore_result == None:
+                print("this time i skip")
+                # self._update_stats(stats, batch_size, loader)
+                continue
+            # get inputs
+            # print(data)
+
+            self.data_to_gpu_time = time.time()
+
+            data['epoch'] = self.epoch
+            data['settings'] = self.settings
+
+            stats = {}
+            reward_record = []
+            miou_record = []
+            e_miou_record = []
+            num_seq = len(data['num_frames'])
+
+            # Calculate reward tensor
+            # reward_tensor = torch.zeros(explore_result['baseline_iou'].size())
+            baseline_iou = explore_result['baseline_iou']
+            # explore_iou = explore_result['explore_iou']
+            for seq_idx in range(num_seq):
+                num_frames = data['num_frames'][seq_idx] - 1
+                b_miou = torch.mean(baseline_iou[:num_frames, seq_idx])
+                #    e_miou = torch.mean(explore_iou[:num_frames, seq_idx])
+                miou_record.append(b_miou.item())
+                #    e_miou_record.append(e_miou.item())
+
+                b_reward = b_miou.item()
+            #    e_reward = e_miou.item()
+            #    iou_gap = e_reward - b_reward
+            #    reward_record.append(iou_gap)
+            #    reward_tensor[:num_frames, seq_idx] = iou_gap
+
+            # Training mode
+            cursor = 0
+            bs_backward = 1
+
+            # print(self.actor.net.module.box_head.decoder.layers[2].mlpx.fc1.weight)
+            self.optimizer.zero_grad()
+            while cursor < num_seq:
+                # print("now is ", cursor , "and all is ", num_seq)
+                model_inputs = {}
+                model_inputs['slt_loss_weight'] = 15
+                if cursor < num_seq:
+                    model_inputs['template_images'] = explore_result['template_images'][
+                                                      cursor:cursor + bs_backward].cuda()
+                else:
+                    model_inputs['template_images'] = explore_result['template_images_reverse'][
+                                                      cursor - num_seq:cursor - num_seq + bs_backward].cuda()
+                model_inputs['search_images'] = explore_result['search_images'][:, cursor:cursor + bs_backward].cuda()
+                model_inputs['search_anno'] = explore_result['search_anno'][:, cursor:cursor + bs_backward].cuda()
+                model_inputs['pre_seq'] = explore_result['pre_seq'][:, cursor:cursor + bs_backward].cuda()
+                model_inputs['x_feat'] = explore_result['x_feat'].squeeze(1)[:, cursor:cursor + bs_backward].cuda()
+                model_inputs['epoch'] = data['epoch']
+                # model_inputs['template_update'] = explore_result['template_update'].squeeze(1)[:,
+                #                                  cursor:cursor + bs_backward].cuda()
+                # print("this is cursor")
+                # print(explore_result['pre_seq'].shape)
+                # print(explore_result['x_feat'].squeeze(1).shape)
+                # model_inputs['action_tensor'] = explore_result['action_tensor'][:, cursor:cursor + bs_backward].cuda()
+                # model_inputs['reward_tensor'] = reward_tensor[:, cursor:cursor + bs_backward].cuda()
+
+                loss, stats_cur = self.actor.compute_sequence_losses(model_inputs)
+                # for name, param in self.actor.net.named_parameters():
+                #    shape, c = (param.grad.shape, param.grad.sum()) if param.grad is not None else (None, None)
+                #    print(f'{name}: {param.shape} \n\t grad: {shape} \n\t {c}')
+                # print("i make this!")
+                loss.backward()
+                # print("i made that?")
+
+                for key, val in stats_cur.items():
+                    if key in stats:
+                        stats[key] += val * (bs_backward / num_seq)
+                    else:
+                        stats[key] = val * (bs_backward / num_seq)
+                cursor += bs_backward
+            grad_norm = clip_grad_norm_(self.actor.net.parameters(), 100)
+            stats['grad_norm'] = grad_norm
+            # print(self.actor.net.module.backbone.blocks[8].mlp.fc1.weight)
+            self.optimizer.step()
+            # print(self.optimizer)
+
+            miou = np.mean(miou_record)
+            self.miou_list.append(miou)
+            # stats['reward'] = np.mean(reward_record)
+            # stats['e_mIoU'] = np.mean(e_miou_record)
+            stats['mIoU'] = miou
+            stats['mIoU10'] = np.mean(self.miou_list[-10:])
+            stats['mIoU100'] = np.mean(self.miou_list[-100:])
+
+            batch_size = num_seq * np.max(data['num_frames'])
+            self._update_stats(stats, batch_size, loader)
+            self._print_stats(i, loader, batch_size)
+            torch.cuda.empty_cache()
+
+            # # forward pass
+            # if not self.use_amp:
+            #     loss, stats = self.actor(data)
+            # else:
+            #     with autocast():
+            #         loss, stats = self.actor(data)
+            #
+            # # backward pass and update weights
+            # if loader.training:
+            #     self.optimizer.zero_grad()
+            #     if not self.use_amp:
+            #         loss.backward()
+            #         if self.settings.grad_clip_norm > 0:
+            #             torch.nn.utils.clip_grad_norm_(self.actor.net.parameters(), self.settings.grad_clip_norm)
+            #         self.optimizer.step()
+            #     else:
+            #         self.scaler.scale(loss).backward()
+            #         self.scaler.step(self.optimizer)
+            #         self.scaler.update()
+
+            # update statistics
+            # batch_size = data['template_images'].shape[loader.stack_dim]
+            # self._update_stats(stats, batch_size, loader)
+
+            # print statistics
+            # self._print_stats(i, loader, batch_size)
+
+            # update wandb status
+            # if self.wandb_writer is not None and i % self.settings.print_interval == 0:
+            #    if self.settings.local_rank in [-1, 0]:
+            #        self.wandb_writer.write_log(self.stats, self.epoch)
+
+        # calculate ETA after every epoch
+        # epoch_time = self.prev_time - self.start_time
+        # print("Epoch Time: " + str(datetime.timedelta(seconds=epoch_time)))
+        # print("Avg Data Time: %.5f" % (self.avg_date_time / self.num_frames * batch_size))
+        # print("Avg GPU Trans Time: %.5f" % (self.avg_gpu_trans_time / self.num_frames * batch_size))
+        # print("Avg Forward Time: %.5f" % (self.avg_forward_time / self.num_frames * batch_size))
+
+    def train_epoch(self):
+        """Do one epoch for each loader."""
+        for loader in self.loaders:
+            if self.epoch % loader.epoch_interval == 0:
+                # 2021.1.10 Set epoch
+                if isinstance(loader.sampler, DistributedSampler):
+                    loader.sampler.set_epoch(self.epoch)
+                self.cycle_dataset(loader)
+
+        self._stats_new_epoch()
+        # if self.settings.local_rank in [-1, 0]:
+        #    self._write_tensorboard()
+
+    def _init_timing(self):
+        self.num_frames = 0
+        self.start_time = time.time()
+        self.prev_time = self.start_time
+        self.avg_date_time = 0
+        self.avg_gpu_trans_time = 0
+        self.avg_forward_time = 0
+
+    def _update_stats(self, new_stats: OrderedDict, batch_size, loader):
+        # Initialize stats if not initialized yet
+        if loader.name not in self.stats.keys() or self.stats[loader.name] is None:
+            self.stats[loader.name] = OrderedDict({name: AverageMeter() for name in new_stats.keys()})
+
+        # add lr state
+        if loader.training:
+            lr_list = self.lr_scheduler.get_last_lr()
+            for i, lr in enumerate(lr_list):
+                var_name = 'LearningRate/group{}'.format(i)
+                if var_name not in self.stats[loader.name].keys():
+                    self.stats[loader.name][var_name] = StatValue()
+                self.stats[loader.name][var_name].update(lr)
+
+        for name, val in new_stats.items():
+            if name not in self.stats[loader.name].keys():
+                self.stats[loader.name][name] = AverageMeter()
+            self.stats[loader.name][name].update(val, batch_size)
+
+    def _print_stats(self, i, loader, batch_size):
+        self.num_frames += batch_size
+        current_time = time.time()
+        batch_fps = batch_size / (current_time - self.prev_time)
+        average_fps = self.num_frames / (current_time - self.start_time)
+        prev_frame_time_backup = self.prev_time
+        self.prev_time = current_time
+
+        self.avg_date_time += (self.data_read_done_time - prev_frame_time_backup)
+        self.avg_gpu_trans_time += (self.data_to_gpu_time - self.data_read_done_time)
+        self.avg_forward_time += current_time - self.data_to_gpu_time
+
+        if i % self.settings.print_interval == 0 or i == loader.__len__():
+            print_str = '[%s: %d, %d / %d] ' % (loader.name, self.epoch, i, loader.__len__())
+            print_str += 'FPS: %.1f (%.1f)  ,  ' % (average_fps, batch_fps)
+
+            # 2021.12.14 add data time print
+            print_str += 'DataTime: %.3f (%.3f)  ,  ' % (
+            self.avg_date_time / self.num_frames * batch_size, self.avg_gpu_trans_time / self.num_frames * batch_size)
+            print_str += 'ForwardTime: %.3f  ,  ' % (self.avg_forward_time / self.num_frames * batch_size)
+            print_str += 'TotalTime: %.3f  ,  ' % ((current_time - self.start_time) / self.num_frames * batch_size)
+            # print_str += 'DataTime: %.3f (%.3f)  ,  ' % (self.data_read_done_time - prev_frame_time_backup, self.data_to_gpu_time - self.data_read_done_time)
+            # print_str += 'ForwardTime: %.3f  ,  ' % (current_time - self.data_to_gpu_time)
+            # print_str += 'TotalTime: %.3f  ,  ' % (current_time - prev_frame_time_backup)
+
+            for name, val in self.stats[loader.name].items():
+                if (self.settings.print_stats is None or name in self.settings.print_stats):
+                    if hasattr(val, 'avg'):
+                        print_str += '%s: %.5f  ,  ' % (name, val.avg)
+                    # else:
+                    #     print_str += '%s: %r  ,  ' % (name, val)
+
+            print(print_str[:-5])
+            log_str = print_str[:-5] + '\n'
+            with open(self.settings.log_file, 'a') as f:
+                f.write(log_str)
+
+    def _stats_new_epoch(self):
+        # Record learning rate
+        for loader in self.loaders:
+            if loader.training:
+                try:
+                    lr_list = self.lr_scheduler.get_last_lr()
+                except:
+                    lr_list = self.lr_scheduler._get_lr(self.epoch)
+                for i, lr in enumerate(lr_list):
+                    var_name = 'LearningRate/group{}'.format(i)
+                    if var_name not in self.stats[loader.name].keys():
+                        self.stats[loader.name][var_name] = StatValue()
+                    self.stats[loader.name][var_name].update(lr)
+
+        for loader_stats in self.stats.values():
+            if loader_stats is None:
+                continue
+            for stat_value in loader_stats.values():
+                if hasattr(stat_value, 'new_epoch'):
+                    stat_value.new_epoch()
+
+    # def _write_tensorboard(self):
+    #    if self.epoch == 1:
+    #        self.tensorboard_writer.write_info(self.settings.script_name, self.settings.description)
+
+    #    self.tensorboard_writer.write_epoch(self.stats, self.epoch)

lib/train/trainers/ltr_trainer.py

+import os
+import datetime
+from collections import OrderedDict
+
+#from lib.train.data.wandb_logger import WandbWriter
+from lib.train.trainers import BaseTrainer
+from lib.train.admin import AverageMeter, StatValue
+#from lib.train.admin import TensorboardWriter
+import torch
+import time
+from torch.utils.data.distributed import DistributedSampler
+from torch.cuda.amp import autocast
+from torch.cuda.amp import GradScaler
+
+from lib.utils.misc import get_world_size
+
+
+class LTRTrainer(BaseTrainer):
+    def __init__(self, actor, loaders, optimizer, settings, lr_scheduler=None, use_amp=False):
+        """
+        args:
+            actor - The actor for training the network
+            loaders - list of dataset loaders, e.g. [train_loader, val_loader]. In each epoch, the trainer runs one
+                        epoch for each loader.
+            optimizer - The optimizer used for training, e.g. Adam
+            settings - Training settings
+            lr_scheduler - Learning rate scheduler
+        """
+        super().__init__(actor, loaders, optimizer, settings, lr_scheduler)
+
+        self._set_default_settings()
+
+        # Initialize statistics variables
+        self.stats = OrderedDict({loader.name: None for loader in self.loaders})
+
+        # Initialize tensorboard and wandb
+        #self.wandb_writer = None
+        #if settings.local_rank in [-1, 0]:
+        #    tensorboard_writer_dir = os.path.join(self.settings.env.tensorboard_dir, self.settings.project_path)
+        #    if not os.path.exists(tensorboard_writer_dir):
+        #        os.makedirs(tensorboard_writer_dir)
+        #    self.tensorboard_writer = TensorboardWriter(tensorboard_writer_dir, [l.name for l in loaders])
+
+        #    if settings.use_wandb:
+        #        world_size = get_world_size()
+        #        cur_train_samples = self.loaders[0].dataset.samples_per_epoch * max(0, self.epoch - 1)
+        #        interval = (world_size * settings.batchsize)  # * interval
+        #        self.wandb_writer = WandbWriter(settings.project_path[6:], {}, tensorboard_writer_dir, cur_train_samples, interval)
+
+        self.move_data_to_gpu = getattr(settings, 'move_data_to_gpu', True)
+        print("move_data", self.move_data_to_gpu)
+        self.settings = settings
+        self.use_amp = use_amp
+        if use_amp:
+            self.scaler = GradScaler()
+
+    def _set_default_settings(self):
+        # Dict of all default values
+        default = {'print_interval': 10,
+                   'print_stats': None,
+                   'description': ''}
+
+        for param, default_value in default.items():
+            if getattr(self.settings, param, None) is None:
+                setattr(self.settings, param, default_value)
+
+    def cycle_dataset(self, loader):
+        """Do a cycle of training or validation."""
+
+        self.actor.train(loader.training)
+        torch.set_grad_enabled(loader.training)
+
+        self._init_timing()
+
+        for i, data in enumerate(loader, 1):
+            self.data_read_done_time = time.time()
+            # get inputs
+            if self.move_data_to_gpu:
+                data = data.to(self.device)
+
+            self.data_to_gpu_time = time.time()
+
+            data['epoch'] = self.epoch
+            data['settings'] = self.settings
+            # forward pass
+            if not self.use_amp:
+                loss, stats = self.actor(data)
+            else:
+                with autocast():
+                    loss, stats = self.actor(data)
+
+            # backward pass and update weights
+            if loader.training:
+                self.optimizer.zero_grad()
+                if not self.use_amp:
+                    loss.backward()
+                    if self.settings.grad_clip_norm > 0:
+                        torch.nn.utils.clip_grad_norm_(self.actor.net.parameters(), self.settings.grad_clip_norm)
+                    self.optimizer.step()
+                else:
+                    self.scaler.scale(loss).backward()
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+
+            # update statistics
+            batch_size = data['template_images'].shape[loader.stack_dim]
+            self._update_stats(stats, batch_size, loader)
+
+            # print statistics
+            self._print_stats(i, loader, batch_size)
+
+            # update wandb status
+            #if self.wandb_writer is not None and i % self.settings.print_interval == 0:
+            #    if self.settings.local_rank in [-1, 0]:
+            #        self.wandb_writer.write_log(self.stats, self.epoch)
+
+        # calculate ETA after every epoch
+        epoch_time = self.prev_time - self.start_time
+        print("Epoch Time: " + str(datetime.timedelta(seconds=epoch_time)))
+        print("Avg Data Time: %.5f" % (self.avg_date_time / self.num_frames * batch_size))
+        print("Avg GPU Trans Time: %.5f" % (self.avg_gpu_trans_time / self.num_frames * batch_size))
+        print("Avg Forward Time: %.5f" % (self.avg_forward_time / self.num_frames * batch_size))
+
+    def train_epoch(self):
+        """Do one epoch for each loader."""
+        for loader in self.loaders:
+            if self.epoch % loader.epoch_interval == 0:
+                # 2021.1.10 Set epoch
+                if isinstance(loader.sampler, DistributedSampler):
+                    loader.sampler.set_epoch(self.epoch)
+                self.cycle_dataset(loader)
+
+        self._stats_new_epoch()
+        #if self.settings.local_rank in [-1, 0]:
+        #    self._write_tensorboard()
+
+    def _init_timing(self):
+        self.num_frames = 0
+        self.start_time = time.time()
+        self.prev_time = self.start_time
+        self.avg_date_time = 0
+        self.avg_gpu_trans_time = 0
+        self.avg_forward_time = 0
+
+    def _update_stats(self, new_stats: OrderedDict, batch_size, loader):
+        # Initialize stats if not initialized yet
+        if loader.name not in self.stats.keys() or self.stats[loader.name] is None:
+            self.stats[loader.name] = OrderedDict({name: AverageMeter() for name in new_stats.keys()})
+
+        # add lr state
+        if loader.training:
+            lr_list = self.lr_scheduler.get_last_lr()
+            for i, lr in enumerate(lr_list):
+                var_name = 'LearningRate/group{}'.format(i)
+                if var_name not in self.stats[loader.name].keys():
+                    self.stats[loader.name][var_name] = StatValue()
+                self.stats[loader.name][var_name].update(lr)
+
+        for name, val in new_stats.items():
+            if name not in self.stats[loader.name].keys():
+                self.stats[loader.name][name] = AverageMeter()
+            self.stats[loader.name][name].update(val, batch_size)
+
+    def _print_stats(self, i, loader, batch_size):
+        self.num_frames += batch_size
+        current_time = time.time()
+        batch_fps = batch_size / (current_time - self.prev_time)
+        average_fps = self.num_frames / (current_time - self.start_time)
+        prev_frame_time_backup = self.prev_time
+        self.prev_time = current_time
+
+        self.avg_date_time += (self.data_read_done_time - prev_frame_time_backup)
+        self.avg_gpu_trans_time += (self.data_to_gpu_time - self.data_read_done_time)
+        self.avg_forward_time += current_time - self.data_to_gpu_time
+
+        if i % self.settings.print_interval == 0 or i == loader.__len__():
+            print_str = '[%s: %d, %d / %d] ' % (loader.name, self.epoch, i, loader.__len__())
+            print_str += 'FPS: %.1f (%.1f)  ,  ' % (average_fps, batch_fps)
+
+            # 2021.12.14 add data time print
+            print_str += 'DataTime: %.3f (%.3f)  ,  ' % (self.avg_date_time / self.num_frames * batch_size, self.avg_gpu_trans_time / self.num_frames * batch_size)
+            print_str += 'ForwardTime: %.3f  ,  ' % (self.avg_forward_time / self.num_frames * batch_size)
+            print_str += 'TotalTime: %.3f  ,  ' % ((current_time - self.start_time) / self.num_frames * batch_size)
+            # print_str += 'DataTime: %.3f (%.3f)  ,  ' % (self.data_read_done_time - prev_frame_time_backup, self.data_to_gpu_time - self.data_read_done_time)
+            # print_str += 'ForwardTime: %.3f  ,  ' % (current_time - self.data_to_gpu_time)
+            # print_str += 'TotalTime: %.3f  ,  ' % (current_time - prev_frame_time_backup)
+
+            for name, val in self.stats[loader.name].items():
+                if (self.settings.print_stats is None or name in self.settings.print_stats):
+                    if hasattr(val, 'avg'):
+                        print_str += '%s: %.5f  ,  ' % (name, val.avg)
+                    # else:
+                    #     print_str += '%s: %r  ,  ' % (name, val)
+
+            print(print_str[:-5])
+            log_str = print_str[:-5] + '\n'
+            with open(self.settings.log_file, 'a') as f:
+                f.write(log_str)
+
+    def _stats_new_epoch(self):
+        # Record learning rate
+        for loader in self.loaders:
+            if loader.training:
+                try:
+                    lr_list = self.lr_scheduler.get_last_lr()
+                except:
+                    lr_list = self.lr_scheduler._get_lr(self.epoch)
+                for i, lr in enumerate(lr_list):
+                    var_name = 'LearningRate/group{}'.format(i)
+                    if var_name not in self.stats[loader.name].keys():
+                        self.stats[loader.name][var_name] = StatValue()
+                    self.stats[loader.name][var_name].update(lr)
+
+        for loader_stats in self.stats.values():
+            if loader_stats is None:
+                continue
+            for stat_value in loader_stats.values():
+                if hasattr(stat_value, 'new_epoch'):
+                    stat_value.new_epoch()
+
+    #def _write_tensorboard(self):
+    #    if self.epoch == 1:
+    #        self.tensorboard_writer.write_info(self.settings.script_name, self.settings.description)
+
+    #    self.tensorboard_writer.write_epoch(self.stats, self.epoch)

lib/utils/__init__.py

+from .tensor import TensorDict, TensorList

lib/utils/box_ops.py

+import torch
+from torchvision.ops.boxes import box_area
+import numpy as np
+
+
+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 box_xywh_to_xyxy(x):
+    x1, y1, w, h = x.unbind(-1)
+    b = [x1, y1, x1 + w, y1 + h]
+    return torch.stack(b, dim=-1)
+
+
+def box_xyxy_to_xywh(x):
+    x1, y1, x2, y2 = x.unbind(-1)
+    b = [x1, y1, x2 - x1, y2 - y1]
+    return torch.stack(b, dim=-1)
+
+
+def box_xyxy_to_cxcywh(x):
+    x0, y0, x1, y1 = x.unbind(-1)
+    b = [(x0 + x1) / 2, (y0 + y1) / 2,
+         (x1 - x0), (y1 - y0)]
+    return torch.stack(b, dim=-1)
+
+
+# modified from torchvision to also return the union
+'''Note that this function only supports shape (N,4)'''
+
+
+def box_iou(boxes1, boxes2):
+    """
+
+    :param boxes1: (N, 4) (x1,y1,x2,y2)
+    :param boxes2: (N, 4) (x1,y1,x2,y2)
+    :return:
+    """
+    area1 = box_area(boxes1) # (N,)
+    area2 = box_area(boxes2) # (N,)
+
+    lt = torch.max(boxes1[:, :2], boxes2[:, :2])  # (N,2)
+    rb = torch.min(boxes1[:, 2:], boxes2[:, 2:])  # (N,2)
+
+    wh = (rb - lt).clamp(min=0)  # (N,2)
+    inter = wh[:, 0] * wh[:, 1]  # (N,)
+
+    union = area1 + area2 - inter
+
+    iou = inter / union
+    return iou, union
+
+
+'''Note that this implementation is different from DETR's'''
+
+
+def generalized_box_iou(boxes1, boxes2):
+    """
+    Generalized IoU from https://giou.stanford.edu/
+
+    The boxes should be in [x0, y0, x1, y1] format
+
+    boxes1: (N, 4)
+    boxes2: (N, 4)
+    """
+    # degenerate boxes gives inf / nan results
+    # so do an early check
+    # try:
+    #assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
+    # assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
+    iou, union = box_iou(boxes1, boxes2) # (N,)
+
+    lt = torch.min(boxes1[:, :2], boxes2[:, :2])
+    rb = torch.max(boxes1[:, 2:], boxes2[:, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # (N,2)
+    area = wh[:, 0] * wh[:, 1] # (N,)
+
+    return iou - (area - union) / area, iou
+
+
+def giou_loss(boxes1, boxes2):
+    """
+
+    :param boxes1: (N, 4) (x1,y1,x2,y2)
+    :param boxes2: (N, 4) (x1,y1,x2,y2)
+    :return:
+    """
+    giou, iou = generalized_box_iou(boxes1, boxes2)
+    return (1 - giou).mean(), iou
+
+
+def clip_box(box: list, H, W, margin=0):
+    x1, y1, w, h = box
+    x2, y2 = x1 + w, y1 + h
+    x1 = min(max(0, x1), W-margin)
+    x2 = min(max(margin, x2), W)
+    y1 = min(max(0, y1), H-margin)
+    y2 = min(max(margin, y2), H)
+    w = max(margin, x2-x1)
+    h = max(margin, y2-y1)
+    return [x1, y1, w, h]

lib/utils/ce_utils.py

+import math
+
+import torch
+import torch.nn.functional as F
+
+
+def generate_bbox_mask(bbox_mask, bbox):
+    b, h, w = bbox_mask.shape
+    for i in range(b):
+        bbox_i = bbox[i].cpu().tolist()
+        bbox_mask[i, int(bbox_i[1]):int(bbox_i[1] + bbox_i[3] - 1), int(bbox_i[0]):int(bbox_i[0] + bbox_i[2] - 1)] = 1
+    return bbox_mask
+
+
+def generate_mask_cond(cfg, bs, device, gt_bbox):
+    template_size = cfg.DATA.TEMPLATE.SIZE
+    stride = cfg.MODEL.BACKBONE.STRIDE
+    template_feat_size = template_size // stride
+
+    if cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'ALL':
+        box_mask_z = None
+    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_POINT':
+        if template_feat_size == 8:
+            index = slice(3, 4)
+        elif template_feat_size == 12:
+            index = slice(5, 6)
+        elif template_feat_size == 7:
+            index = slice(3, 4)
+        elif template_feat_size == 14:
+            index = slice(6, 7)
+        else:
+            raise NotImplementedError
+        box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
+        box_mask_z[:, index, index] = 1
+        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
+    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'CTR_REC':
+        # use fixed 4x4 region, 3:5 for 8x8
+        # use fixed 4x4 region 5:6 for 12x12
+        if template_feat_size == 8:
+            index = slice(3, 5)
+        elif template_feat_size == 12:
+            index = slice(5, 7)
+        elif template_feat_size == 7:
+            index = slice(3, 4)
+        else:
+            raise NotImplementedError
+        box_mask_z = torch.zeros([bs, template_feat_size, template_feat_size], device=device)
+        box_mask_z[:, index, index] = 1
+        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
+
+    elif cfg.MODEL.BACKBONE.CE_TEMPLATE_RANGE == 'GT_BOX':
+        box_mask_z = torch.zeros([bs, template_size, template_size], device=device)
+        # box_mask_z_ori = data['template_seg'][0].view(-1, 1, *data['template_seg'].shape[2:])  # (batch, 1, 128, 128)
+        box_mask_z = generate_bbox_mask(box_mask_z, gt_bbox * template_size).unsqueeze(1).to(
+            torch.float)  # (batch, 1, 128, 128)
+        # box_mask_z_vis = box_mask_z.cpu().numpy()
+        box_mask_z = F.interpolate(box_mask_z, scale_factor=1. / cfg.MODEL.BACKBONE.STRIDE, mode='bilinear',
+                                   align_corners=False)
+        box_mask_z = box_mask_z.flatten(1).to(torch.bool)
+        # box_mask_z_vis = box_mask_z[:, 0, ...].cpu().numpy()
+        # gaussian_maps_vis = generate_heatmap(data['template_anno'], self.cfg.DATA.TEMPLATE.SIZE, self.cfg.MODEL.STRIDE)[0].cpu().numpy()
+    else:
+        raise NotImplementedError
+
+    return box_mask_z
+
+
+def adjust_keep_rate(epoch, warmup_epochs, total_epochs, ITERS_PER_EPOCH, base_keep_rate=0.5, max_keep_rate=1, iters=-1):
+    if epoch < warmup_epochs:
+        return 1
+    if epoch >= total_epochs:
+        return base_keep_rate
+    if iters == -1:
+        iters = epoch * ITERS_PER_EPOCH
+    total_iters = ITERS_PER_EPOCH * (total_epochs - warmup_epochs)
+    iters = iters - ITERS_PER_EPOCH * warmup_epochs
+    keep_rate = base_keep_rate + (max_keep_rate - base_keep_rate) \
+        * (math.cos(iters / total_iters * math.pi) + 1) * 0.5
+
+    return keep_rate

lib/utils/focal_loss.py

+from abc import ABC
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FocalLoss(nn.Module, ABC):
+    def __init__(self, alpha=2, beta=4):
+        super(FocalLoss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+
+    def forward(self, prediction, target):
+        positive_index = target.eq(1).float()
+        negative_index = target.lt(1).float()
+
+        negative_weights = torch.pow(1 - target, self.beta)
+        # clamp min value is set to 1e-12 to maintain the numerical stability
+        prediction = torch.clamp(prediction, 1e-12)
+
+        positive_loss = torch.log(prediction) * torch.pow(1 - prediction, self.alpha) * positive_index
+        negative_loss = torch.log(1 - prediction) * torch.pow(prediction,
+                                                              self.alpha) * negative_weights * negative_index
+
+        num_positive = positive_index.float().sum()
+        positive_loss = positive_loss.sum()
+        negative_loss = negative_loss.sum()
+
+        if num_positive == 0:
+            loss = -negative_loss
+        else:
+            loss = -(positive_loss + negative_loss) / num_positive
+
+        return loss
+
+
+class LBHinge(nn.Module):
+    """Loss that uses a 'hinge' on the lower bound.
+    This means that for samples with a label value smaller than the threshold, the loss is zero if the prediction is
+    also smaller than that threshold.
+    args:
+        error_matric:  What base loss to use (MSE by default).
+        threshold:  Threshold to use for the hinge.
+        clip:  Clip the loss if it is above this value.
+    """
+    def __init__(self, error_metric=nn.MSELoss(), threshold=None, clip=None):
+        super().__init__()
+        self.error_metric = error_metric
+        self.threshold = threshold if threshold is not None else -100
+        self.clip = clip
+
+    def forward(self, prediction, label, target_bb=None):
+        negative_mask = (label < self.threshold).float()
+        positive_mask = (1.0 - negative_mask)
+
+        prediction = negative_mask * F.relu(prediction) + positive_mask * prediction
+
+        loss = self.error_metric(prediction, positive_mask * label)
+
+        if self.clip is not None:
+            loss = torch.min(loss, torch.tensor([self.clip], device=loss.device))
+        return loss
\ No newline at end of file

lib/utils/heapmap_utils.py

+import numpy as np
+import torch
+
+
+def generate_heatmap(bboxes, patch_size=320, stride=16):
+    """
+    Generate ground truth heatmap same as CenterNet
+    Args:
+        bboxes (torch.Tensor): shape of [num_search, bs, 4]
+
+    Returns:
+        gaussian_maps: list of generated heatmap
+
+    """
+    gaussian_maps = []
+    heatmap_size = patch_size // stride
+    for single_patch_bboxes in bboxes:
+        bs = single_patch_bboxes.shape[0]
+        gt_scoremap = torch.zeros(bs, heatmap_size, heatmap_size)
+        classes = torch.arange(bs).to(torch.long)
+        bbox = single_patch_bboxes * heatmap_size
+        wh = bbox[:, 2:]
+        centers_int = (bbox[:, :2] + wh / 2).round()
+        CenterNetHeatMap.generate_score_map(gt_scoremap, classes, wh, centers_int, 0.7)
+        gaussian_maps.append(gt_scoremap.to(bbox.device))
+    return gaussian_maps
+
+
+class CenterNetHeatMap(object):
+    @staticmethod
+    def generate_score_map(fmap, gt_class, gt_wh, centers_int, min_overlap):
+        radius = CenterNetHeatMap.get_gaussian_radius(gt_wh, min_overlap)
+        radius = torch.clamp_min(radius, 0)
+        radius = radius.type(torch.int).cpu().numpy()
+        for i in range(gt_class.shape[0]):
+            channel_index = gt_class[i]
+            CenterNetHeatMap.draw_gaussian(fmap[channel_index], centers_int[i], radius[i])
+
+    @staticmethod
+    def get_gaussian_radius(box_size, min_overlap):
+        """
+        copyed from CornerNet
+        box_size (w, h), it could be a torch.Tensor, numpy.ndarray, list or tuple
+        notice: we are using a bug-version, please refer to fix bug version in CornerNet
+        """
+        # box_tensor = torch.Tensor(box_size)
+        box_tensor = box_size
+        width, height = box_tensor[..., 0], box_tensor[..., 1]
+
+        a1 = 1
+        b1 = height + width
+        c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
+        sq1 = torch.sqrt(b1 ** 2 - 4 * a1 * c1)
+        r1 = (b1 + sq1) / 2
+
+        a2 = 4
+        b2 = 2 * (height + width)
+        c2 = (1 - min_overlap) * width * height
+        sq2 = torch.sqrt(b2 ** 2 - 4 * a2 * c2)
+        r2 = (b2 + sq2) / 2
+
+        a3 = 4 * min_overlap
+        b3 = -2 * min_overlap * (height + width)
+        c3 = (min_overlap - 1) * width * height
+        sq3 = torch.sqrt(b3 ** 2 - 4 * a3 * c3)
+        r3 = (b3 + sq3) / 2
+
+        return torch.min(r1, torch.min(r2, r3))
+
+    @staticmethod
+    def gaussian2D(radius, sigma=1):
+        # m, n = [(s - 1.) / 2. for s in shape]
+        m, n = radius
+        y, x = np.ogrid[-m: m + 1, -n: n + 1]
+
+        gauss = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
+        gauss[gauss < np.finfo(gauss.dtype).eps * gauss.max()] = 0
+        return gauss
+
+    @staticmethod
+    def draw_gaussian(fmap, center, radius, k=1):
+        diameter = 2 * radius + 1
+        gaussian = CenterNetHeatMap.gaussian2D((radius, radius), sigma=diameter / 6)
+        gaussian = torch.Tensor(gaussian)
+        x, y = int(center[0]), int(center[1])
+        height, width = fmap.shape[:2]
+
+        left, right = min(x, radius), min(width - x, radius + 1)
+        top, bottom = min(y, radius), min(height - y, radius + 1)
+
+        masked_fmap = fmap[y - top: y + bottom, x - left: x + right]
+        masked_gaussian = gaussian[radius - top: radius + bottom, radius - left: radius + right]
+        if min(masked_gaussian.shape) > 0 and min(masked_fmap.shape) > 0:
+            masked_fmap = torch.max(masked_fmap, masked_gaussian * k)
+            fmap[y - top: y + bottom, x - left: x + right] = masked_fmap
+        # return fmap
+
+
+def compute_grids(features, strides):
+    """
+    grids regret to the input image size
+    """
+    grids = []
+    for level, feature in enumerate(features):
+        h, w = feature.size()[-2:]
+        shifts_x = torch.arange(
+            0, w * strides[level],
+            step=strides[level],
+            dtype=torch.float32, device=feature.device)
+        shifts_y = torch.arange(
+            0, h * strides[level],
+            step=strides[level],
+            dtype=torch.float32, device=feature.device)
+        shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
+        shift_x = shift_x.reshape(-1)
+        shift_y = shift_y.reshape(-1)
+        grids_per_level = torch.stack((shift_x, shift_y), dim=1) + \
+                          strides[level] // 2
+        grids.append(grids_per_level)
+    return grids
+
+
+def get_center3x3(locations, centers, strides, range=3):
+    '''
+    Inputs:
+        locations: M x 2
+        centers: N x 2
+        strides: M
+    '''
+    range = (range - 1) / 2
+    M, N = locations.shape[0], centers.shape[0]
+    locations_expanded = locations.view(M, 1, 2).expand(M, N, 2)  # M x N x 2
+    centers_expanded = centers.view(1, N, 2).expand(M, N, 2)  # M x N x 2
+    strides_expanded = strides.view(M, 1, 1).expand(M, N, 2)  # M x N
+    centers_discret = ((centers_expanded / strides_expanded).int() * strides_expanded).float() + \
+                      strides_expanded / 2  # M x N x 2
+    dist_x = (locations_expanded[:, :, 0] - centers_discret[:, :, 0]).abs()
+    dist_y = (locations_expanded[:, :, 1] - centers_discret[:, :, 1]).abs()
+    return (dist_x <= strides_expanded[:, :, 0] * range) & \
+           (dist_y <= strides_expanded[:, :, 0] * range)
+
+
+def get_pred(score_map_ctr, size_map, offset_map, feat_size):
+    max_score, idx = torch.max(score_map_ctr.flatten(1), dim=1, keepdim=True)
+
+    idx = idx.unsqueeze(1).expand(idx.shape[0], 2, 1)
+    size = size_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+    offset = offset_map.flatten(2).gather(dim=2, index=idx).squeeze(-1)
+
+    return size * feat_size, offset

lib/utils/lmdb_utils.py

+import lmdb
+import numpy as np
+import cv2
+import json
+
+LMDB_ENVS = dict()
+LMDB_HANDLES = dict()
+LMDB_FILELISTS = dict()
+
+
+def get_lmdb_handle(name):
+    global LMDB_HANDLES, LMDB_FILELISTS
+    item = LMDB_HANDLES.get(name, None)
+    if item is None:
+        env = lmdb.open(name, readonly=True, lock=False, readahead=False, meminit=False)
+        LMDB_ENVS[name] = env
+        item = env.begin(write=False)
+        LMDB_HANDLES[name] = item
+
+    return item
+
+
+def decode_img(lmdb_fname, key_name):
+    handle = get_lmdb_handle(lmdb_fname)
+    binfile = handle.get(key_name.encode())
+    if binfile is None:
+        print("Illegal data detected. %s %s" % (lmdb_fname, key_name))
+    s = np.frombuffer(binfile, np.uint8)
+    x = cv2.cvtColor(cv2.imdecode(s, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
+    return x
+
+
+def decode_str(lmdb_fname, key_name):
+    handle = get_lmdb_handle(lmdb_fname)
+    binfile = handle.get(key_name.encode())
+    string = binfile.decode()
+    return string
+
+
+def decode_json(lmdb_fname, key_name):
+    return json.loads(decode_str(lmdb_fname, key_name))
+
+
+if __name__ == "__main__":
+    lmdb_fname = "/data/sda/v-yanbi/iccv21/LittleBoy_clean/data/got10k_lmdb"
+    '''Decode image'''
+    # key_name = "test/GOT-10k_Test_000001/00000001.jpg"
+    # img = decode_img(lmdb_fname, key_name)
+    # cv2.imwrite("001.jpg", img)
+    '''Decode str'''
+    # key_name = "test/list.txt"
+    # key_name = "train/GOT-10k_Train_000001/groundtruth.txt"
+    key_name = "train/GOT-10k_Train_000001/absence.label"
+    str_ = decode_str(lmdb_fname, key_name)
+    print(str_)

lib/utils/merge.py

+import torch
+
+
+def merge_template_search(inp_list, return_search=False, return_template=False):
+    """NOTICE: search region related features must be in the last place"""
+    seq_dict = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0),
+                "mask": torch.cat([x["mask"] for x in inp_list], dim=1),
+                "pos": torch.cat([x["pos"] for x in inp_list], dim=0)}
+    if return_search:
+        x = inp_list[-1]
+        seq_dict.update({"feat_x": x["feat"], "mask_x": x["mask"], "pos_x": x["pos"]})
+    if return_template:
+        z = inp_list[0]
+        seq_dict.update({"feat_z": z["feat"], "mask_z": z["mask"], "pos_z": z["pos"]})
+    return seq_dict
+
+
+def get_qkv(inp_list):
+    """The 1st element of the inp_list is about the template,
+    the 2nd (the last) element is about the search region"""
+    dict_x = inp_list[-1]
+    dict_c = {"feat": torch.cat([x["feat"] for x in inp_list], dim=0),
+              "mask": torch.cat([x["mask"] for x in inp_list], dim=1),
+              "pos": torch.cat([x["pos"] for x in inp_list], dim=0)}  # concatenated dict
+    q = dict_x["feat"] + dict_x["pos"]
+    k = dict_c["feat"] + dict_c["pos"]
+    v = dict_c["feat"]
+    key_padding_mask = dict_c["mask"]
+    return q, k, v, key_padding_mask

lib/utils/misc.py

+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Misc functions, including distributed helpers.
+
+Mostly copy-paste from torchvision references.
+"""
+import os
+import subprocess
+import time
+from collections import defaultdict, deque
+import datetime
+import pickle
+from typing import Optional, List
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+# needed due to empty tensor bug in pytorch and torchvision 0.5
+import torchvision
+vers = torchvision.__version__.split('.')
+if int(vers[0]) <= 0 and int(vers[1]) < 7:
+    from torchvision.ops import _new_empty_tensor
+    from torchvision.ops.misc import _output_size
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+
+    # serialized to a Tensor
+    buffer = pickle.dumps(data)
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to("cuda")
+
+    # obtain Tensor size of each rank
+    local_size = torch.tensor([tensor.numel()], device="cuda")
+    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]
+    dist.all_gather(size_list, local_size)
+    size_list = [int(size.item()) for size in size_list]
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    tensor_list = []
+    for _ in size_list:
+        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))
+    if local_size != max_size:
+        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")
+        tensor = torch.cat((tensor, padding), dim=0)
+    dist.all_gather(tensor_list, tensor)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}',
+                'max mem: {memory:.0f}'
+            ])
+        else:
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}'
+            ])
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
+    sha = 'N/A'
+    diff = "clean"
+    branch = 'N/A'
+    try:
+        sha = _run(['git', 'rev-parse', 'HEAD'])
+        subprocess.check_output(['git', 'diff'], cwd=cwd)
+        diff = _run(['git', 'diff-index', 'HEAD'])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def collate_fn(batch):
+    batch = list(zip(*batch))
+    batch[0] = nested_tensor_from_tensor_list(batch[0])
+    return tuple(batch)
+
+
+def _max_by_axis(the_list):
+    # type: (List[List[int]]) -> List[int]
+    maxes = the_list[0] # get the first one
+    for sublist in the_list[1:]: # [h,w,3]
+        for index, item in enumerate(sublist): # index: 0,1,2
+            maxes[index] = max(maxes[index], item) # compare current max with the other elements in the whole
+    return maxes
+
+
+class NestedTensor(object):
+    def __init__(self, tensors, mask: Optional[Tensor]):
+        self.tensors = tensors
+        self.mask = mask
+
+    def to(self, device):
+        # type: (Device) -> NestedTensor # noqa
+        cast_tensor = self.tensors.to(device)
+        mask = self.mask
+        if mask is not None:
+            assert mask is not None
+            cast_mask = mask.to(device)
+        else:
+            cast_mask = None
+        return NestedTensor(cast_tensor, cast_mask)
+
+    def decompose(self):
+        return self.tensors, self.mask
+
+    def __repr__(self):
+        return str(self.tensors)
+
+
+def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
+    # TODO make this more general
+    if tensor_list[0].ndim == 3:
+        if torchvision._is_tracing():
+            # nested_tensor_from_tensor_list() does not export well to ONNX
+            # call _onnx_nested_tensor_from_tensor_list() instead
+            return _onnx_nested_tensor_from_tensor_list(tensor_list)
+
+        # TODO make it support different-sized images
+        max_size = _max_by_axis([list(img.shape) for img in tensor_list]) # [[3,h1,w1], [3,h2,w2], [3,h3,w3], ...]
+        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
+        batch_shape = [len(tensor_list)] + max_size # ()
+        b, c, h, w = batch_shape
+        dtype = tensor_list[0].dtype
+        device = tensor_list[0].device
+        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
+        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
+        for img, pad_img, m in zip(tensor_list, tensor, mask):
+            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) # copy valid regions of the images to the largest padded base.
+            m[: img.shape[1], :img.shape[2]] = False
+    else:
+        raise ValueError('not supported')
+    return NestedTensor(tensor, mask)
+
+
+# _onnx_nested_tensor_from_tensor_list() is an implementation of
+# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
+@torch.jit.unused
+def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
+    max_size = []
+    for i in range(tensor_list[0].dim()):
+        max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64)
+        max_size.append(max_size_i)
+    max_size = tuple(max_size)
+
+    # work around for
+    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+    # m[: img.shape[1], :img.shape[2]] = False
+    # which is not yet supported in onnx
+    padded_imgs = []
+    padded_masks = []
+    for img in tensor_list:
+        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
+        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
+        padded_imgs.append(padded_img)
+
+        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
+        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
+        padded_masks.append(padded_mask.to(torch.bool))
+
+    tensor = torch.stack(padded_imgs)
+    mask = torch.stack(padded_masks)
+
+    return NestedTensor(tensor, mask=mask)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}'.format(
+        args.rank, args.dist_url), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+@torch.no_grad()
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    if target.numel() == 0:
+        return [torch.zeros([], device=output.device)]
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
+    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
+    """
+    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
+    This will eventually be supported natively by PyTorch, and this
+    class can go away.
+    """
+    if float(torchvision.__version__[:3]) < 0.7:
+        if input.numel() > 0:
+            return torch.nn.functional.interpolate(
+                input, size, scale_factor, mode, align_corners
+            )
+
+        output_shape = _output_size(2, input, size, scale_factor)
+        output_shape = list(input.shape[:-2]) + list(output_shape)
+        return _new_empty_tensor(input, output_shape)
+    else:
+        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)

lib/utils/tensor.py

+import functools
+import torch
+import copy
+from collections import OrderedDict
+
+
+class TensorDict(OrderedDict):
+    """Container mainly used for dicts of torch tensors. Extends OrderedDict with pytorch functionality."""
+
+    def concat(self, other):
+        """Concatenates two dicts without copying internal data."""
+        return TensorDict(self, **other)
+
+    def copy(self):
+        return TensorDict(super(TensorDict, self).copy())
+
+    def __deepcopy__(self, memodict={}):
+        return TensorDict(copy.deepcopy(list(self), memodict))
+
+    def __getattr__(self, name):
+        if not hasattr(torch.Tensor, name):
+            raise AttributeError('\'TensorDict\' object has not attribute \'{}\''.format(name))
+
+        def apply_attr(*args, **kwargs):
+            return TensorDict({n: getattr(e, name)(*args, **kwargs) if hasattr(e, name) else e for n, e in self.items()})
+        return apply_attr
+
+    def attribute(self, attr: str, *args):
+        return TensorDict({n: getattr(e, attr, *args) for n, e in self.items()})
+
+    def apply(self, fn, *args, **kwargs):
+        return TensorDict({n: fn(e, *args, **kwargs) for n, e in self.items()})
+
+    @staticmethod
+    def _iterable(a):
+        return isinstance(a, (TensorDict, list))
+
+
+class TensorList(list):
+    """Container mainly used for lists of torch tensors. Extends lists with pytorch functionality."""
+
+    def __init__(self, list_of_tensors = None):
+        if list_of_tensors is None:
+            list_of_tensors = list()
+        super(TensorList, self).__init__(list_of_tensors)
+
+    def __deepcopy__(self, memodict={}):
+        return TensorList(copy.deepcopy(list(self), memodict))
+
+    def __getitem__(self, item):
+        if isinstance(item, int):
+            return super(TensorList, self).__getitem__(item)
+        elif isinstance(item, (tuple, list)):
+            return TensorList([super(TensorList, self).__getitem__(i) for i in item])
+        else:
+            return TensorList(super(TensorList, self).__getitem__(item))
+
+    def __add__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 + e2 for e1, e2 in zip(self, other)])
+        return TensorList([e + other for e in self])
+
+    def __radd__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 + e1 for e1, e2 in zip(self, other)])
+        return TensorList([other + e for e in self])
+
+    def __iadd__(self, other):
+        if TensorList._iterable(other):
+            for i, e2 in enumerate(other):
+                self[i] += e2
+        else:
+            for i in range(len(self)):
+                self[i] += other
+        return self
+
+    def __sub__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 - e2 for e1, e2 in zip(self, other)])
+        return TensorList([e - other for e in self])
+
+    def __rsub__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 - e1 for e1, e2 in zip(self, other)])
+        return TensorList([other - e for e in self])
+
+    def __isub__(self, other):
+        if TensorList._iterable(other):
+            for i, e2 in enumerate(other):
+                self[i] -= e2
+        else:
+            for i in range(len(self)):
+                self[i] -= other
+        return self
+
+    def __mul__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 * e2 for e1, e2 in zip(self, other)])
+        return TensorList([e * other for e in self])
+
+    def __rmul__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 * e1 for e1, e2 in zip(self, other)])
+        return TensorList([other * e for e in self])
+
+    def __imul__(self, other):
+        if TensorList._iterable(other):
+            for i, e2 in enumerate(other):
+                self[i] *= e2
+        else:
+            for i in range(len(self)):
+                self[i] *= other
+        return self
+
+    def __truediv__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 / e2 for e1, e2 in zip(self, other)])
+        return TensorList([e / other for e in self])
+
+    def __rtruediv__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 / e1 for e1, e2 in zip(self, other)])
+        return TensorList([other / e for e in self])
+
+    def __itruediv__(self, other):
+        if TensorList._iterable(other):
+            for i, e2 in enumerate(other):
+                self[i] /= e2
+        else:
+            for i in range(len(self)):
+                self[i] /= other
+        return self
+
+    def __matmul__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 @ e2 for e1, e2 in zip(self, other)])
+        return TensorList([e @ other for e in self])
+
+    def __rmatmul__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 @ e1 for e1, e2 in zip(self, other)])
+        return TensorList([other @ e for e in self])
+
+    def __imatmul__(self, other):
+        if TensorList._iterable(other):
+            for i, e2 in enumerate(other):
+                self[i] @= e2
+        else:
+            for i in range(len(self)):
+                self[i] @= other
+        return self
+
+    def __mod__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 % e2 for e1, e2 in zip(self, other)])
+        return TensorList([e % other for e in self])
+
+    def __rmod__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e2 % e1 for e1, e2 in zip(self, other)])
+        return TensorList([other % e for e in self])
+
+    def __pos__(self):
+        return TensorList([+e for e in self])
+
+    def __neg__(self):
+        return TensorList([-e for e in self])
+
+    def __le__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 <= e2 for e1, e2 in zip(self, other)])
+        return TensorList([e <= other for e in self])
+
+    def __ge__(self, other):
+        if TensorList._iterable(other):
+            return TensorList([e1 >= e2 for e1, e2 in zip(self, other)])
+        return TensorList([e >= other for e in self])
+
+    def concat(self, other):
+        return TensorList(super(TensorList, self).__add__(other))
+
+    def copy(self):
+        return TensorList(super(TensorList, self).copy())
+
+    def unroll(self):
+        if not any(isinstance(t, TensorList) for t in self):
+            return self
+
+        new_list = TensorList()
+        for t in self:
+            if isinstance(t, TensorList):
+                new_list.extend(t.unroll())
+            else:
+                new_list.append(t)
+        return new_list
+
+    def list(self):
+        return list(self)
+
+    def attribute(self, attr: str, *args):
+        return TensorList([getattr(e, attr, *args) for e in self])
+
+    def apply(self, fn):
+        return TensorList([fn(e) for e in self])
+
+    def __getattr__(self, name):
+        if not hasattr(torch.Tensor, name):
+            raise AttributeError('\'TensorList\' object has not attribute \'{}\''.format(name))
+
+        def apply_attr(*args, **kwargs):
+            return TensorList([getattr(e, name)(*args, **kwargs) for e in self])
+
+        return apply_attr
+
+    @staticmethod
+    def _iterable(a):
+        return isinstance(a, (TensorList, list))
+
+
+def tensor_operation(op):
+    def islist(a):
+        return isinstance(a, TensorList)
+
+    @functools.wraps(op)
+    def oplist(*args, **kwargs):
+        if len(args) == 0:
+            raise ValueError('Must be at least one argument without keyword (i.e. operand).')
+
+        if len(args) == 1:
+            if islist(args[0]):
+                return TensorList([op(a, **kwargs) for a in args[0]])
+        else:
+            # Multiple operands, assume max two
+            if islist(args[0]) and islist(args[1]):
+                return TensorList([op(a, b, *args[2:], **kwargs) for a, b in zip(*args[:2])])
+            if islist(args[0]):
+                return TensorList([op(a, *args[1:], **kwargs) for a in args[0]])
+            if islist(args[1]):
+                return TensorList([op(args[0], b, *args[2:], **kwargs) for b in args[1]])
+
+        # None of the operands are lists
+        return op(*args, **kwargs)
+
+    return oplist

lib/utils/variable_hook.py

+import torch
+from bytecode import Bytecode, Instr
+
+
+class get_local(object):
+    cache = {}
+    is_activate = False
+
+    def __init__(self, varname):
+        self.varname = varname
+
+    def __call__(self, func):
+        if not type(self).is_activate:
+            return func
+
+        type(self).cache[func.__qualname__] = []
+        c = Bytecode.from_code(func.__code__)
+        extra_code = [
+            Instr('STORE_FAST', '_res'),
+            Instr('LOAD_FAST', self.varname),
+            Instr('STORE_FAST', '_value'),
+            Instr('LOAD_FAST', '_res'),
+            Instr('LOAD_FAST', '_value'),
+            Instr('BUILD_TUPLE', 2),
+            Instr('STORE_FAST', '_result_tuple'),
+            Instr('LOAD_FAST', '_result_tuple'),
+        ]
+        c[-1:-1] = extra_code
+        func.__code__ = c.to_code()
+
+        def wrapper(*args, **kwargs):
+            res, values = func(*args, **kwargs)
+            if isinstance(values, torch.Tensor):
+                type(self).cache[func.__qualname__].append(values.detach().cpu().numpy())
+            elif isinstance(values, list):  # list of Tensor
+                type(self).cache[func.__qualname__].append([value.detach().cpu().numpy() for value in values])
+            else:
+                raise NotImplementedError
+            return res
+
+        return wrapper
+
+    @classmethod
+    def clear(cls):
+        for key in cls.cache.keys():
+            cls.cache[key] = []
+
+    @classmethod
+    def activate(cls):
+        cls.is_activate = True

model.properties

+# 模型唯一标识
+modelCode = 1449
+# 模型名称
+modelName=samurai_pytorch
+# 模型描述
+modelDescription=基于SAM2的Zero-Shot视觉跟踪模型
+# 应用场景
+appScenario=AIGC,零售,制造,电商,医疗,教育
+# 框架类型
+frameType=pytorch

requirements.txt

+matplotlib==3.7
+tikzplotlib 
+jpeg4py 
+opencv-python 
+lmdb 
+pandas 
+scipy 
+loguru
+
+numpy==1.24.4
\ No newline at end of file

sam2/.clang-format

+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands:   false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+  AfterClass:      false
+  AfterControlStatement: false
+  AfterEnum:       false
+  AfterFunction:   false
+  AfterNamespace:  false
+  AfterObjCDeclaration: false
+  AfterStruct:     false
+  AfterUnion:      false
+  BeforeCatch:     false
+  BeforeElse:      false
+  IndentBraces:    false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit:     80
+CommentPragmas:  '^ IWYU pragma:'
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat:   false
+ForEachMacros:   [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ]
+IncludeCategories:
+  - Regex:           '^<.*\.h(pp)?>'
+    Priority:        1
+  - Regex:           '^<.*'
+    Priority:        2
+  - Regex:           '.*'
+    Priority:        3
+IndentCaseLabels: true
+IndentWidth:     2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ''
+MacroBlockEnd:   ''
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 200
+PointerAlignment: Left
+ReflowComments:  true
+SortIncludes:    true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles:  false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard:        Cpp11
+TabWidth:        8
+UseTab:          Never