def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) dist = False synchronize() # init logging file logger = Logger(cfg.work_dir / "log_test.txt") sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build model model = build_model( cfg, 0) # use num_classes=0 since we do not need classifier for testing model.cuda() if dist: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[cfg.gpu], output_device=cfg.gpu, find_unused_parameters=True, ) # elif cfg.total_gpus > 1: else: model = torch.nn.DataParallel(model) # load checkpoint state_dict = load_checkpoint(args.resume) # load test data_loader test_loaders, queries, galleries = build_test_dataloader(cfg) for key in state_dict: if not key.startswith("state_dict"): continue print("==> Test with {}".format(key)) copy_state_dict(state_dict[key], model) # start testing for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid(cfg, model, loader, query, gallery, dataset_name=cfg.TEST.datasets[i]) # print time end_time = time.monotonic() print("Total running time: ", timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) set_random_seed(cfg.TRAIN.seed, cfg.TRAIN.deterministic) synchronize() # init logging file logger = Logger(cfg.work_dir / "log.txt", debug=False) sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, train_sets = build_train_dataloader(cfg) # the number of classes for the model is tricky, # you need to make sure that # it is always larger than the number of clusters num_classes = 0 for idx, set in enumerate(train_sets): if idx in cfg.TRAIN.unsup_dataset_indexes: # number of clusters in an unsupervised dataset # must not be larger than the number of images num_classes += len(set) else: # ground-truth classes for supervised dataset num_classes += set.num_pids # build model no.1 model_1 = build_model(cfg, num_classes, init=cfg.MODEL.source_pretrained) model_1.cuda() # build model no.2 model_2 = build_model(cfg, num_classes, init=cfg.MODEL.source_pretrained) model_2.cuda() if dist: ddp_cfg = { "device_ids": [cfg.gpu], "output_device": cfg.gpu, "find_unused_parameters": True, } model_1 = torch.nn.parallel.DistributedDataParallel(model_1, **ddp_cfg) model_2 = torch.nn.parallel.DistributedDataParallel(model_2, **ddp_cfg) elif cfg.total_gpus > 1: model_1 = torch.nn.DataParallel(model_1) model_2 = torch.nn.DataParallel(model_2) # build optimizer optimizer = build_optimizer([model_1, model_2], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions criterions = build_loss(cfg.TRAIN.LOSS, num_classes=num_classes, cuda=True) # build runner runner = MMTRunner( cfg, [ model_1, model_2, ], optimizer, criterions, train_loader, train_sets=train_sets, lr_scheduler=lr_scheduler, reset_optim=True, ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / "model_best.pth") # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): for idx in range(len(runner.model)): print("==> Test on the no.{} model".format(idx)) # test_reid() on self.model[idx] will only evaluate the 'mean_net' # for testing 'net', use self.model[idx].module.net cmc, mAP = test_reid( cfg, runner.model[idx], loader, query, gallery, dataset_name=cfg.TEST.datasets[i], ) # print time end_time = time.monotonic() print("Total running time: ", timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) synchronize() # init logging file logger = Logger(cfg.work_dir / 'log.txt', debug=False) sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, train_sets = build_train_dataloader(cfg) # the number of classes for the model is tricky, # you need to make sure that # it is always larger than the number of clusters num_classes = 0 for idx, set in enumerate(train_sets): if (idx in cfg.TRAIN.unsup_dataset_indexes): # number of clusters in an unsupervised dataset # must not be larger than the number of images num_classes += len(set) else: # ground-truth classes for supervised dataset num_classes += set.num_pids # build model model = build_model(cfg, num_classes, init=cfg.MODEL.source_pretrained) model.cuda() if dist: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[cfg.gpu], output_device=cfg.gpu, find_unused_parameters=True, ) elif (cfg.total_gpus > 1): model = torch.nn.DataParallel(model) # build optimizer optimizer = build_optimizer([ model, ], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions criterions = build_loss(cfg.TRAIN, num_classes=num_classes, cuda=True) # build runner runner = BaseRunner( cfg, model, optimizer, criterions, train_loader, train_sets=train_sets, lr_scheduler=lr_scheduler, reset_optim=True, ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / 'model_best.pth') # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid(cfg, model, loader, query, gallery, dataset_name=cfg.TEST.datasets[i]) # print time end_time = time.monotonic() print('Total running time: ', timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) synchronize() # init logging file logger = Logger(cfg.work_dir / 'log.txt', debug=False) sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, train_sets = build_train_dataloader(cfg) # build model model = build_model(cfg, 0, init=cfg.MODEL.source_pretrained) model.cuda() if dist: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[cfg.gpu], output_device=cfg.gpu, find_unused_parameters=True, ) elif (cfg.total_gpus>1): model = torch.nn.DataParallel(model) # build optimizer optimizer = build_optimizer([model,], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions criterions = build_loss(cfg.TRAIN.LOSS, triplet_key='feat', cuda=True) # build runner runner = BaseRunner( cfg, model, optimizer, criterions, train_loader, train_sets = train_sets, lr_scheduler = lr_scheduler, meter_formats = {'Time': ':.3f'}, reset_optim = False, ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / 'model_best.pth') # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid( cfg, model, loader, query, gallery, dataset_name = cfg.TEST.datasets[i] ) # print time end_time = time.monotonic() print('Total running time: ', timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) synchronize() # init logging file logger = Logger(cfg.work_dir / "log.txt", debug=False) sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, train_sets = build_train_dataloader(cfg, joint=False) # build model model = build_model(cfg, 0, init=cfg.MODEL.source_pretrained) model.cuda() if dist: ddp_cfg = { "device_ids": [cfg.gpu], "output_device": cfg.gpu, "find_unused_parameters": True, } model = DistributedDataParallel(model, **ddp_cfg) elif cfg.total_gpus > 1: model = DataParallel(model) # build optimizer optimizer = build_optimizer([model], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions num_memory = 0 for idx, set in enumerate(train_sets): if idx in cfg.TRAIN.unsup_dataset_indexes: # instance-level memory for unlabeled data num_memory += len(set) else: # class-level memory for labeled data num_memory += set.num_pids if isinstance(model, (DataParallel, DistributedDataParallel)): num_features = model.module.num_features else: num_features = model.num_features criterions = build_loss( cfg.TRAIN.LOSS, num_features=num_features, num_memory=num_memory, cuda=True, ) # init memory loaders, datasets = build_val_dataloader( cfg, for_clustering=True, all_datasets=True ) memory_features = [] for idx, (loader, dataset) in enumerate(zip(loaders, datasets)): features = extract_features( model, loader, dataset, with_path=False, prefix="Extract: ", ) assert features.size(0) == len(dataset) if idx in cfg.TRAIN.unsup_dataset_indexes: # init memory for unlabeled data with instance features memory_features.append(features) else: # init memory for labeled data with class centers centers_dict = collections.defaultdict(list) for i, (_, pid, _) in enumerate(dataset): centers_dict[pid].append(features[i].unsqueeze(0)) centers = [ torch.cat(centers_dict[pid], 0).mean(0) for pid in sorted(centers_dict.keys()) ] memory_features.append(torch.stack(centers, 0)) del loaders, datasets memory_features = torch.cat(memory_features) criterions["hybrid_memory"]._update_feature(memory_features) # build runner runner = SpCLRunner( cfg, model, optimizer, criterions, train_loader, train_sets=train_sets, lr_scheduler=lr_scheduler, meter_formats={"Time": ":.3f",}, reset_optim=False, ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / "model_best.pth") # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid( cfg, model, loader, query, gallery, dataset_name=cfg.TEST.datasets[i] ) # print time end_time = time.monotonic() print("Total running time: ", timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) set_random_seed(cfg.TRAIN.seed, cfg.TRAIN.deterministic) synchronize() # init logging file logger = Logger(cfg.work_dir / "log.txt", debug=False) sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, train_sets = build_train_dataloader(cfg, joint=False) # build model model = build_model(cfg, 0, init=cfg.MODEL.source_pretrained) model.cuda() if dist: ddp_cfg = { "device_ids": [cfg.gpu], "output_device": cfg.gpu, "find_unused_parameters": True, } model = DistributedDataParallel(model, **ddp_cfg) elif cfg.total_gpus > 1: model = DataParallel(model) # build optimizer optimizer = build_optimizer([model], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions num_memory = 0 for idx, set in enumerate(train_sets): if idx in cfg.TRAIN.unsup_dataset_indexes: # instance-level memory for unlabeled data num_memory += len(set) else: # class-level memory for labeled data num_memory += set.num_pids if isinstance(model, (DataParallel, DistributedDataParallel)): num_features = model.module.num_features else: num_features = model.num_features criterions = build_loss( cfg.TRAIN.LOSS, num_features=num_features, num_memory=num_memory, cuda=True, ) # build runner runner = SSBRunner( cfg, model, optimizer, criterions, train_loader, train_sets=train_sets, lr_scheduler=lr_scheduler, meter_formats={ "Time": ":.3f", }, reset_optim=False, ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / "model_best.pth") # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid(cfg, model, loader, query, gallery, dataset_name=cfg.TEST.datasets[i]) # print time end_time = time.monotonic() print("Total running time: ", timedelta(seconds=end_time - start_time))
def main(): start_time = time.monotonic() # init distributed training args, cfg = parge_config() dist = init_dist(cfg) set_random_seed(cfg.TRAIN.seed, cfg.TRAIN.deterministic) synchronize() # init logging file logger = Logger(cfg.work_dir / "log.txt") sys.stdout = logger print("==========\nArgs:{}\n==========".format(args)) log_config_to_file(cfg) # build train loader train_loader, _ = build_train_dataloader(cfg) # build model model = build_model(cfg, train_loader.loader.dataset.num_pids) model.cuda() if dist: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[cfg.gpu], output_device=cfg.gpu, find_unused_parameters=True, ) elif cfg.total_gpus > 1: model = torch.nn.DataParallel(model) # build optimizer optimizer = build_optimizer([model], **cfg.TRAIN.OPTIM) # build lr_scheduler if cfg.TRAIN.SCHEDULER.lr_scheduler is not None: lr_scheduler = build_lr_scheduler(optimizer, **cfg.TRAIN.SCHEDULER) else: lr_scheduler = None # build loss functions criterions = build_loss( cfg.TRAIN.LOSS, num_classes=train_loader.loader.dataset.num_pids, cuda=True ) # build runner runner = BaseRunner( cfg, model, optimizer, criterions, train_loader, lr_scheduler=lr_scheduler ) # resume if args.resume_from: runner.resume(args.resume_from) # start training runner.run() # load the best model runner.resume(cfg.work_dir / "model_best.pth") # final testing test_loaders, queries, galleries = build_test_dataloader(cfg) for i, (loader, query, gallery) in enumerate(zip(test_loaders, queries, galleries)): cmc, mAP = test_reid( cfg, model, loader, query, gallery, dataset_name=cfg.TEST.datasets[i] ) # print time end_time = time.monotonic() print("Total running time: ", timedelta(seconds=end_time - start_time))