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))
