def update_sampler(sampler, model, loader, query, gallery, sub_set, vlad=True, gpu=None, sync_gather=False):
if (dist.get_rank()==0):
print ("===> Start extracting features for sorting gallery")
features = extract_features(model, loader, sorted(list(set(query) | set(gallery))),
vlad=vlad, gpu=gpu, sync_gather=sync_gather)
distmat, _, _ = pairwise_distance(features, query, gallery)
del features
if (dist.get_rank()==0):
print ("===> Start sorting gallery")
sampler.sort_gallery(distmat, sub_set)
del distmat
def update_sampler(sampler, model, loader, query, gallery, sub_set, rerank=False,
vlad=True, gpu=None, sync_gather=False, lambda_value=0.1):
if (dist.get_rank()==0):
print ("===> Start extracting features for sorting gallery")
features = extract_features(model, loader, sorted(list(set(query) | set(gallery))),
vlad=vlad, gpu=gpu, sync_gather=sync_gather)
distmat, _, _ = pairwise_distance(features, query, gallery)
if rerank:
distmat_qq, _, _ = pairwise_distance(features, query, query)
distmat_gg, _, _ = pairwise_distance(features, gallery, gallery)
distmat_jac = re_ranking(distmat.numpy(), distmat_qq.numpy(), distmat_gg.numpy(),
k1=20, k2=1, lambda_value=lambda_value)
distmat_jac = torch.from_numpy(distmat_jac)
del distmat_qq, distmat_gg
else:
distmat_jac = distmat
del features
if (dist.get_rank()==0):
print ("===> Start sorting gallery")
sampler.sort_gallery(distmat, distmat_jac, sub_set)
del distmat, distmat_jac
def main_worker(args):
global start_epoch, best_recall5
init_dist(args.launcher, args)
synchronize()
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
if args.deterministic:
cudnn.deterministic = True
cudnn.benchmark = False
print("Use GPU: {} for training, rank no.{} of world_size {}"
.format(args.gpu, args.rank, args.world_size))
if (args.rank==0):
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
iters = args.iters if (args.iters>0) else None
dataset, train_loader, val_loader, test_loader, sampler, train_extract_loader = get_data(args, iters)
# Create model
model = get_model(args)
# Load from checkpoint
if args.resume:
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model)
start_epoch = checkpoint['epoch']+1
best_recall5 = checkpoint['best_recall5']
if (args.rank==0):
print("=> Start epoch {} best recall5 {:.1%}"
.format(start_epoch, best_recall5))
# Evaluator
evaluator = Evaluator(model)
if (args.rank==0):
print("Test the initial model:")
recalls = evaluator.evaluate(val_loader, sorted(list(set(dataset.q_val) | set(dataset.db_val))),
dataset.q_val, dataset.db_val, dataset.val_pos,
vlad=args.vlad, gpu=args.gpu, sync_gather=args.sync_gather)
# Optimizer
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()),
lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=0.5)
# Trainer
trainer = Trainer(model, margin=args.margin**0.5, gpu=args.gpu)
if ((args.cache_size<args.tuple_size) or (args.cache_size>len(dataset.q_train))):
args.cache_size = len(dataset.q_train)
# Start training
for epoch in range(start_epoch, args.epochs):
sampler.set_epoch(args.seed+epoch)
args.cache_size = args.cache_size * (2 ** (epoch // args.step_size))
g = torch.Generator()
g.manual_seed(args.seed+epoch)
subset_indices = torch.randperm(len(dataset.q_train), generator=g).long().split(args.cache_size)
for subid, subset in enumerate(subset_indices):
update_sampler(sampler, model, train_extract_loader, dataset.q_train, dataset.db_train, subset.tolist(),
vlad=args.vlad, gpu=args.gpu, sync_gather=args.sync_gather)
synchronize()
trainer.train(epoch, subid, train_loader, optimizer,
train_iters=len(train_loader), print_freq=args.print_freq,
vlad=args.vlad, loss_type=args.loss_type)
synchronize()
if ((epoch+1)%args.eval_step==0 or (epoch==args.epochs-1)):
recalls = evaluator.evaluate(val_loader, sorted(list(set(dataset.q_val) | set(dataset.db_val))),
dataset.q_val, dataset.db_val, dataset.val_pos,
vlad=args.vlad, gpu=args.gpu, sync_gather=args.sync_gather)
is_best = recalls[1] > best_recall5
best_recall5 = max(recalls[1], best_recall5)
if (args.rank==0):
save_checkpoint({
'state_dict': model.state_dict(),
'epoch': epoch,
'best_recall5': best_recall5,
}, is_best, fpath=osp.join(args.logs_dir, 'checkpoint'+str(epoch)+'.pth.tar'))
print('\n * Finished epoch {:3d} recall@1: {:5.1%} recall@5: {:5.1%} recall@10: {:5.1%} best@5: {:5.1%}{}\n'.
format(epoch, recalls[0], recalls[1], recalls[2], best_recall5, ' *' if is_best else ''))
lr_scheduler.step()
synchronize()
# final inference
if (args.rank==0):
print("Performing PCA reduction on the best model:")
model.load_state_dict(load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))['state_dict'])
pca_parameters_path = osp.join(args.logs_dir, 'pca_params_model_best.h5')
pca = PCA(args.features, (not args.nowhiten), pca_parameters_path)
dict_f = extract_features(model, train_extract_loader, sorted(list(set(dataset.q_train) | set(dataset.db_train))),
vlad=args.vlad, gpu=args.gpu, sync_gather=args.sync_gather)
features = list(dict_f.values())
if (len(features)>10000):
features = random.sample(features, 10000)
features = torch.stack(features)
if (args.rank==0):
pca.train(features)
synchronize()
del features
if (args.rank==0):
print("Testing on Pitts30k-test:")
evaluator.evaluate(test_loader, sorted(list(set(dataset.q_test) | set(dataset.db_test))),
dataset.q_test, dataset.db_test, dataset.test_pos,
vlad=args.vlad, pca=pca, gpu=args.gpu, sync_gather=args.sync_gather)
synchronize()
return
def main_worker(args):
init_dist(args.launcher, args)
synchronize()
cudnn.benchmark = True
print("Use GPU: {} for testing, rank no.{} of world_size {}".format(
args.gpu, args.rank, args.world_size))
assert (args.resume)
if (args.rank == 0):
log_dir = osp.dirname(args.resume)
sys.stdout = Logger(
osp.join(log_dir, 'log_test_' + args.dataset + '.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
dataset, pitts_train, train_extract_loader, test_loader_q, test_loader_db = get_data(
args)
# Create model
model = get_model(args)
# Load from checkpoint
if args.resume:
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model)
start_epoch = checkpoint['epoch']
best_recall5 = checkpoint['best_recall5']
if (args.rank == 0):
print("=> Start epoch {} best recall5 {:.1%}".format(
start_epoch, best_recall5))
# Evaluator
evaluator = Evaluator(model)
if (args.reduction):
pca_parameters_path = osp.join(
osp.dirname(args.resume),
'pca_params_' + osp.basename(args.resume).split('.')[0] + '.h5')
pca = PCA(args.features, (not args.nowhiten), pca_parameters_path)
if (not osp.isfile(pca_parameters_path)):
dict_f = extract_features(model,
train_extract_loader,
pitts_train,
vlad=args.vlad,
gpu=args.gpu,
sync_gather=args.sync_gather)
features = list(dict_f.values())
if (len(features) > 10000):
features = random.sample(features, 10000)
features = torch.stack(features)
if (args.rank == 0):
pca.train(features)
synchronize()
del features
else:
pca = None
if (args.rank == 0):
print("Evaluate on the test set:")
evaluator.evaluate(test_loader_q,
sorted(list(set(dataset.q_test)
| set(dataset.db_test))),
dataset.q_test,
dataset.db_test,
dataset.test_pos,
gallery_loader=test_loader_db,
vlad=args.vlad,
pca=pca,
rerank=args.rerank,
gpu=args.gpu,
sync_gather=args.sync_gather,
nms=(True if args.dataset == 'tokyo' else False),
rr_topk=args.rr_topk,
lambda_value=args.lambda_value)
synchronize()
return