def build_datamanager(cfg): if cfg.data.type == 'image': return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg)) else: return torchreid.data.VideoDataManager(**videodata_kwargs(cfg))
def main(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument( '--config-file', type=str, default='', help='path to config file' ) parser.add_argument( '-s', '--sources', type=str, nargs='+', help='source datasets (delimited by space)' ) parser.add_argument( '-t', '--targets', type=str, nargs='+', help='target datasets (delimited by space)' ) parser.add_argument( '--transforms', type=str, nargs='+', help='data augmentation' ) parser.add_argument( '--root', type=str, default='', help='path to data root' ) parser.add_argument( '--gpu-devices', type=str, default='', ) parser.add_argument( 'opts', default=None, nargs=argparse.REMAINDER, help='Modify config options using the command-line' ) args = parser.parse_args() cfg = get_default_config() cfg.use_gpu = torch.cuda.is_available() if args.config_file: cfg.merge_from_file(args.config_file) reset_config(cfg, args) cfg.merge_from_list(args.opts) set_random_seed(cfg.train.seed) if cfg.use_gpu and args.gpu_devices: # if gpu_devices is not specified, all available gpus will be used os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices log_name = 'test.log' if cfg.test.evaluate else 'train.log' log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) print('Show configuration\n{}\n'.format(cfg)) print('Collecting env info ...') print('** System info **\n{}\n'.format(collect_env_info())) if cfg.use_gpu: torch.backends.cudnn.benchmark = True datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg)) print('Building model: {}'.format(cfg.model.name)) model = osnet_models.build_model( cfg.model.name, num_classes=datamanager.num_train_pids ) num_params, flops = compute_model_complexity( model, (1, 3, cfg.data.height, cfg.data.width) ) print('Model complexity: params={:,} flops={:,}'.format(num_params, flops)) if cfg.use_gpu: model = nn.DataParallel(model).cuda() optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg)) scheduler = torchreid.optim.build_lr_scheduler( optimizer, **lr_scheduler_kwargs(cfg) ) if cfg.model.resume and check_isfile(cfg.model.resume): cfg.train.start_epoch = resume_from_checkpoint( cfg.model.resume, model, optimizer=optimizer ) print('Building NAS engine') engine = ImageSoftmaxNASEngine( datamanager, model, optimizer, scheduler=scheduler, use_gpu=cfg.use_gpu, label_smooth=cfg.loss.softmax.label_smooth, mc_iter=cfg.nas.mc_iter, init_lmda=cfg.nas.init_lmda, min_lmda=cfg.nas.min_lmda, lmda_decay_step=cfg.nas.lmda_decay_step, lmda_decay_rate=cfg.nas.lmda_decay_rate, fixed_lmda=cfg.nas.fixed_lmda ) engine.run(**engine_run_kwargs(cfg)) print('*** Display the found architecture ***') if cfg.use_gpu: model.module.build_child_graph() else: model.build_child_graph()
def main(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--config-file', type=str, default='', help='path to config file') parser.add_argument( '--gpu-devices', type=str, default='', ) parser.add_argument('opts', default=None, nargs=argparse.REMAINDER, help='Modify config options using the command-line') args = parser.parse_args() cfg = get_default_config() if args.config_file: cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) set_random_seed(cfg.train.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices log_name = 'test.log' if cfg.test.evaluate else 'train.log' log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) print('Show configuration\n{}\n'.format(cfg)) torch.backends.cudnn.benchmark = True datamanager = ImageDataManager(**imagedata_kwargs(cfg)) trainloader, queryloader, galleryloader = datamanager.return_dataloaders() print('Building model: {}'.format(cfg.model.name)) model = build_model(cfg.model.name, datamanager.num_train_pids, 'softmax', pretrained=cfg.model.pretrained) if cfg.model.load_weights and check_isfile(cfg.model.load_weights): load_pretrained_weights(model, cfg.model.load_weights) model = nn.DataParallel(model).cuda() criterion = CrossEntropyLoss(datamanager.num_train_pids, label_smooth=cfg.loss.softmax.label_smooth) optimizer = build_optimizer(model, **optimizer_kwargs(cfg)) scheduler = build_lr_scheduler(optimizer, **lr_scheduler_kwargs(cfg)) if cfg.model.resume and check_isfile(cfg.model.resume): cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume, model, optimizer=optimizer) if cfg.test.evaluate: distmat = evaluate(model, queryloader, galleryloader, dist_metric=cfg.test.dist_metric, normalize_feature=cfg.test.normalize_feature, rerank=cfg.test.rerank, return_distmat=True) if cfg.test.visrank: visualize_ranked_results(distmat, datamanager.return_testdataset(), 'image', width=cfg.data.width, height=cfg.data.height, save_dir=osp.join(cfg.data.save_dir, 'visrank')) return time_start = time.time() print('=> Start training') for epoch in range(cfg.train.start_epoch, cfg.train.max_epoch): train(epoch, cfg.train.max_epoch, model, criterion, optimizer, trainloader, fixbase_epoch=cfg.train.fixbase_epoch, open_layers=cfg.train.open_layers) scheduler.step() if (epoch + 1) % cfg.test.eval_freq == 0 or (epoch + 1) == cfg.train.max_epoch: rank1 = evaluate(model, queryloader, galleryloader, dist_metric=cfg.test.dist_metric, normalize_feature=cfg.test.normalize_feature, rerank=cfg.test.rerank) save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': epoch + 1, 'rank1': rank1, 'optimizer': optimizer.state_dict(), }, cfg.data.save_dir) elapsed = round(time.time() - time_start) elapsed = str(datetime.timedelta(seconds=elapsed)) print('Elapsed {}'.format(elapsed))
def main(): parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('--config-file', type=str, default='', help='path to config file') parser.add_argument('-s', '--sources', type=str, nargs='+', help='source datasets (delimited by space)') parser.add_argument('-t', '--targets', type=str, nargs='+', help='target datasets (delimited by space)') parser.add_argument('--transforms', type=str, nargs='+', help='data augmentation') parser.add_argument('--root', type=str, default='', help='path to data root') parser.add_argument('opts', default=None, nargs=argparse.REMAINDER, help='Modify config options using the command-line') args = parser.parse_args() cfg = get_default_config() cfg.use_gpu = torch.cuda.is_available() if args.config_file: cfg.merge_from_file(args.config_file) reset_config(cfg, args) cfg.merge_from_list(args.opts) set_random_seed(cfg.train.seed) log_name = 'test.log' if cfg.test.evaluate else 'train.log' log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name)) print('Show configuration\n{}\n'.format(cfg)) print('Collecting env info ...') print('** System info **\n{}\n'.format(collect_env_info())) if cfg.use_gpu: torch.backends.cudnn.benchmark = True datamanager = torchreid.data.ImageDataManager(**imagedata_kwargs(cfg)) print('Building model-1: {}'.format(cfg.model.name)) model1 = torchreid.models.build_model( name=cfg.model.name, num_classes=datamanager.num_train_pids, loss=cfg.loss.name, pretrained=cfg.model.pretrained, use_gpu=cfg.use_gpu) num_params, flops = compute_model_complexity( model1, (1, 3, cfg.data.height, cfg.data.width)) print('Model complexity: params={:,} flops={:,}'.format(num_params, flops)) print('Copying model-1 to model-2') model2 = copy.deepcopy(model1) if cfg.model.load_weights1 and check_isfile(cfg.model.load_weights1): load_pretrained_weights(model1, cfg.model.load_weights1) if cfg.model.load_weights2 and check_isfile(cfg.model.load_weights2): load_pretrained_weights(model2, cfg.model.load_weights2) if cfg.use_gpu: model1 = nn.DataParallel(model1).cuda() model2 = nn.DataParallel(model2).cuda() optimizer1 = torchreid.optim.build_optimizer(model1, **optimizer_kwargs(cfg)) scheduler1 = torchreid.optim.build_lr_scheduler(optimizer1, **lr_scheduler_kwargs(cfg)) optimizer2 = torchreid.optim.build_optimizer(model2, **optimizer_kwargs(cfg)) scheduler2 = torchreid.optim.build_lr_scheduler(optimizer2, **lr_scheduler_kwargs(cfg)) if cfg.model.resume1 and check_isfile(cfg.model.resume1): cfg.train.start_epoch = resume_from_checkpoint(cfg.model.resume1, model1, optimizer=optimizer1, scheduler=scheduler1) if cfg.model.resume2 and check_isfile(cfg.model.resume2): resume_from_checkpoint(cfg.model.resume2, model2, optimizer=optimizer2, scheduler=scheduler2) print('Building DML-engine for image-reid') engine = ImageDMLEngine(datamanager, model1, optimizer1, scheduler1, model2, optimizer2, scheduler2, margin=cfg.loss.triplet.margin, weight_t=cfg.loss.triplet.weight_t, weight_x=cfg.loss.triplet.weight_x, weight_ml=cfg.loss.dml.weight_ml, use_gpu=cfg.use_gpu, label_smooth=cfg.loss.softmax.label_smooth, deploy=cfg.model.deploy) engine.run(**engine_run_kwargs(cfg))