def main(): # 1. Get input arguments args = get_args() # 2. Create config instance from args above 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 # 3. Create DataManager Instance datamanager = build_datamanager(cfg) print('Building model: {}'.format(cfg.model.name)) model = 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( model, (1, 3, cfg.data.height, cfg.data.width)) print('Model complexity: params={:,} flops={:,}'.format(num_params, flops)) if cfg.model.load_weights and check_isfile(cfg.model.load_weights): load_pretrained_weights(model, cfg.model.load_weights) 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, scheduler=scheduler) print('Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)) # Build engine and run engine = build_engine(cfg, datamanager, model, optimizer, scheduler) engine.run(**engine_run_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 = build_datamanager(cfg) print('Building model: {}'.format(cfg.model.name)) model = 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( model, (1, 3, cfg.data.height, cfg.data.width)) print('Model complexity: params={:,} flops={:,}'.format(num_params, flops)) if cfg.model.load_weights and check_isfile(cfg.model.load_weights): load_pretrained_weights(model, cfg.model.load_weights) 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 {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)) engine = build_engine(cfg, datamanager, model, optimizer, scheduler) engine.run(**engine_run_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( '--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 = build_datamanager(cfg) model_factory = { 'resnet50_fc512': resnet50_fc512, 'osnet_x1_0': osnet_x1_0, # mixstyle models 'resnet50_fc512_ms12_a0d1': resnet50_fc512_ms12_a0d1, 'resnet50_fc512_ms12_a0d2': resnet50_fc512_ms12_a0d2, 'resnet50_fc512_ms12_a0d3': resnet50_fc512_ms12_a0d3, 'resnet50_fc512_ms12_a0d1_domprior': resnet50_fc512_ms12_a0d1_domprior, 'osnet_x1_0_ms23_a0d1': osnet_x1_0_ms23_a0d1, 'osnet_x1_0_ms23_a0d2': osnet_x1_0_ms23_a0d2, 'osnet_x1_0_ms23_a0d3': osnet_x1_0_ms23_a0d3, 'osnet_x1_0_ms23_a0d1_domprior': osnet_x1_0_ms23_a0d1_domprior, # ablation 'resnet50_fc512_ms1_a0d1': resnet50_fc512_ms1_a0d1, 'resnet50_fc512_ms123_a0d1': resnet50_fc512_ms123_a0d1, 'resnet50_fc512_ms1234_a0d1': resnet50_fc512_ms1234_a0d1, 'resnet50_fc512_ms14_a0d1': resnet50_fc512_ms14_a0d1, 'resnet50_fc512_ms23_a0d1': resnet50_fc512_ms23_a0d1, # dropblock models 'resnet50_fc512_db12': resnet50_fc512_db12, 'osnet_x1_0_db23': osnet_x1_0_db23 } print('Building model: {}'.format(cfg.model.name)) model = model_factory[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( model, (1, 3, cfg.data.height, cfg.data.width)) print('Model complexity: params={:,} flops={:,}'.format(num_params, flops)) if cfg.model.load_weights and check_isfile(cfg.model.load_weights): load_pretrained_weights(model, cfg.model.load_weights) 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, scheduler=scheduler) print('Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)) engine = build_engine(cfg, datamanager, model, optimizer, scheduler) engine.run(**engine_run_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('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))