def download_datasets(datasets, dataset_url, globals): """Download data associated with each required dataset""" if "val_eccv20" in datasets: download_train(globals['root_path']) io_helpers.download_files(["retrieval-SfM-120k-val-eccv2020.pkl"], globals['root_path'] / "train/retrieval-SfM-120k", dataset_url, logfunc=globals["logger"].info) elif "train" in datasets: download_train(globals['root_path']) if "roxford5k" in datasets or "rparis6k" in datasets: download_test(globals['root_path'])
def main(): global args, min_loss, global_step args = parser.parse_args() if args.wandb: # initialize wandb wandb.init( project='cnnimageretrieval-pytorch', name=args.directory, entity='ronaldseoh') # save args provided for this experiment to wandb wandb.config.update(args) # manually check if there are unknown test datasets for dataset in args.test_datasets.split(','): if dataset not in test_datasets_names: raise ValueError('Unsupported or unknown test dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.training_dataset) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) if args.local_whitening: directory += "_lwhiten" if args.regional: directory += "_r" if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin) directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay) directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size) directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size) args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set cuda visible device os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # set random seeds torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) # initialize model if args.pretrained: print(">> Using pre-trained model '{}'".format(args.arch)) else: print(">> Using model from scratch (random weights) '{}'".format(args.arch)) model_params = {} model_params['architecture'] = args.arch model_params['pooling'] = args.pool model_params['local_whitening'] = args.local_whitening model_params['regional'] = args.regional model_params['whitening'] = args.whitening # model_params['mean'] = ... # will use default # model_params['std'] = ... # will use default model_params['pretrained'] = args.pretrained model = init_network(model_params) # move network to gpu model.cuda() # define loss function (criterion) and optimizer if args.loss == 'contrastive': criterion = ContrastiveLoss(margin=args.loss_margin).cuda() elif args.loss == 'triplet': criterion = TripletLoss(margin=args.loss_margin).cuda() else: raise(RuntimeError("Loss {} not available!".format(args.loss))) # parameters split into features, pool, whitening # IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM parameters = [] # add feature parameters parameters.append({'params': model.features.parameters()}) # add local whitening if exists if model.lwhiten is not None: parameters.append({'params': model.lwhiten.parameters()}) # add pooling parameters (or regional whitening which is part of the pooling layer!) if not args.regional: # global, only pooling parameter p weight decay should be 0 if args.pool == 'gem': parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0}) elif args.pool == 'gemmp': parameters.append({'params': model.pool.parameters(), 'lr': args.lr*100, 'weight_decay': 0}) else: # regional, pooling parameter p weight decay should be 0, # and we want to add regional whitening if it is there if args.pool == 'gem': parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, 'weight_decay': 0}) elif args.pool == 'gemmp': parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*100, 'weight_decay': 0}) if model.pool.whiten is not None: parameters.append({'params': model.pool.whiten.parameters()}) # add final whitening if exists if model.whiten is not None: parameters.append({'params': model.whiten.parameters()}) # define optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay) # define learning rate decay schedule # TODO: maybe pass as argument in future implementation? exp_decay = math.exp(-0.01) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 if args.resume: args.resume = os.path.join(args.directory, args.resume) if os.path.isfile(args.resume): # load checkpoint weights and update model and optimizer print(">> Loading checkpoint:\n>> '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] min_loss = checkpoint['min_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) # important not to forget scheduler updating scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch']-1) else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading code normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) if args.query_size == -1: train_query_size = float('Inf') else: train_query_size = args.query_size if args.pool_size == -1: train_pool_size = float('Inf') else: train_pool_size = args.pool_size train_dataset = TuplesDataset( name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=train_query_size, poolsize=train_pool_size, transform=transform, store_nidxs_others_up_to=args.store_nidxs_others_up_to, store_nidxs_others_order_by=args.store_nidxs_others_order_by, totally_random_nidxs=args.totally_random_nidxs, totally_random_nidxs_others=args.totally_random_nidxs_others, dense_refresh_batch_and_nearby=args.dense_refresh_batch_and_nearby, dense_refresh_batch_multi_hop=args.dense_refresh_batch_multi_hop, dense_refresh_batch_random=args.dense_refresh_batch_random, ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True, sampler=None, drop_last=True, collate_fn=collate_tuples ) if args.val: val_dataset = TuplesDataset( name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=float('Inf'), poolsize=float('Inf'), transform=transform ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=collate_tuples ) if args.wandb: # Start watching 'model' from wandb wandb.watch(model) # evaluate the network before starting # this might not be necessary? test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=-1) indexes_to_refresh = [] for epoch in range(start_epoch, args.epochs): # set manual seeds per epoch random.seed(args.seed + epoch) np.random.seed(args.seed + epoch) torch.manual_seed(args.seed + epoch) torch.cuda.manual_seed_all(args.seed + epoch) # train for one epoch on train set loss, indexes_to_refresh = train(train_loader, model, criterion, optimizer, epoch, indexes_to_refresh) if args.wandb: wandb.log({"loss_avg": loss, "epoch": epoch, "global_step": global_step}) ## This is average loss # adjust learning rate for each epoch scheduler.step() # # debug printing to check if everything ok # lr_feat = optimizer.param_groups[0]['lr'] # lr_pool = optimizer.param_groups[1]['lr'] # print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool)) # evaluate on validation set if args.val: with torch.no_grad(): loss = validate(val_loader, model, criterion, epoch) if args.wandb: wandb.log({"loss_validation": loss, "epoch": epoch, "global_step": global_step}) # evaluate on test datasets every test_freq epochs if (epoch + 1) % args.test_freq == 0: with torch.no_grad(): test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=epoch) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) if is_best: print("Epoch", str(epoch + 1), "lower loss:", min_loss) save_checkpoint({ 'epoch': epoch + 1, 'meta': model.meta, 'state_dict': model.state_dict(), 'min_loss': min_loss, 'optimizer' : optimizer.state_dict(), }, is_best, args.directory) # calculate avg_neg_distance and avg_pos_distance for one last time print("Training finished. Calculating the final avg_{neg,pos}_distance...") avg_neg_distance, _ = train_loader.dataset.create_epoch_tuples( model, batch_members=[], refresh_query_selection=False, refresh_query_vectors=False, refresh_negative_pool=False, refresh_negative_pool_vectors=False, refresh_nidxs=False, refresh_nidxs_vectors=False) if args.wandb: wandb.log({"avg_neg_distance": avg_neg_distance, 'epoch': epoch, "global_step": global_step}) if args.calculate_positive_distance: avg_pos_distance = train_loader.dataset.calculate_average_positive_distance() if args.wandb: wandb.log({"avg_pos_distance": avg_pos_distance, 'epoch': epoch, "global_step": global_step})
def main(): args = parser.parse_args() # check if there are unknown datasets for dataset in args.datasets.split(','): if dataset not in datasets_names: raise ValueError( 'Unsupported or unknown dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path if args.network_path is not None: print(">> Loading network:\n>>>> '{}'".format(args.network_path)) if args.network_path in PRETRAINED: # pretrained networks (downloaded automatically) state = load_url(PRETRAINED[args.network_path], model_dir=os.path.join(get_data_root(), 'networks')) else: # fine-tuned network from path state = torch.load(args.network_path) # parsing net params from meta # architecture, pooling, mean, std required # the rest has default values, in case that is doesnt exist net_params = {} net_params['architecture'] = state['meta']['architecture'] net_params['pooling'] = state['meta']['pooling'] net_params['local_whitening'] = state['meta'].get( 'local_whitening', False) net_params['regional'] = state['meta'].get('regional', False) net_params['whitening'] = state['meta'].get('whitening', False) net_params['mean'] = state['meta']['mean'] net_params['std'] = state['meta']['std'] net_params['pretrained'] = False # load network net = init_network(net_params) net.load_state_dict(state['state_dict']) # if whitening is precomputed if 'Lw' in state['meta']: net.meta['Lw'] = state['meta']['Lw'] print(">>>> loaded network: ") print(net.meta_repr()) # loading offtheshelf network elif args.network_offtheshelf is not None: # parse off-the-shelf parameters offtheshelf = args.network_offtheshelf.split('-') net_params = {} net_params['architecture'] = offtheshelf[0] net_params['pooling'] = offtheshelf[1] net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:] net_params['regional'] = 'reg' in offtheshelf[2:] net_params['whitening'] = 'whiten' in offtheshelf[2:] net_params['pretrained'] = True # load off-the-shelf network print(">> Loading off-the-shelf network:\n>>>> '{}'".format( args.network_offtheshelf)) net = init_network(net_params) print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[ 'regional'] and not net.meta['whitening']: msp = net.pool.p.item() print(">> Set-up multiscale:") print(">>>> ms: {}".format(ms)) print(">>>> msp: {}".format(msp)) else: msp = 1 # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # compute whitening if args.whitening is not None: start = time.time() if 'Lw' in net.meta and args.whitening in net.meta['Lw']: print('>> {}: Whitening is precomputed, loading it...'.format( args.whitening)) if len(ms) > 1: Lw = net.meta['Lw'][args.whitening]['ms'] else: Lw = net.meta['Lw'][args.whitening]['ss'] else: # if we evaluate networks from path we should save/load whitening # not to compute it every time if args.network_path is not None: whiten_fn = args.network_path + '_{}_whiten'.format( args.whitening) if len(ms) > 1: whiten_fn += '_ms' whiten_fn += '.pth' else: whiten_fn = None if whiten_fn is not None and os.path.isfile(whiten_fn): print('>> {}: Whitening is precomputed, loading it...'.format( args.whitening)) Lw = torch.load(whiten_fn) else: print('>> {}: Learning whitening...'.format(args.whitening)) # loading db db_root = os.path.join(get_data_root(), 'train', args.whitening) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [ cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids'])) ] # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # learning whitening print('>> {}: Learning...'.format(args.whitening)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} # saving whitening if whiten_fn exists if whiten_fn is not None: print('>> {}: Saving to {}...'.format( args.whitening, whiten_fn)) torch.save(Lw, whiten_fn) print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time() - start))) else: Lw = None # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # extract database and query vectors print('>> {}: database images...'.format(dataset)) images = get_imlist("E:\\PycharmProjects\\image-retrieval\\holiday2\\") names = [] for i, img_path in enumerate(images): img_name = os.path.split(img_path)[1] print(img_name) names.append(img_name) # prepare config structure for the test dataset # cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) # images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] # try: # bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] # except: # bbxs = None # for holidaysmanrot and copydays # names = [] # for i, img_path in enumerate(images): # img_name = os.path.split(img_path)[1] # print(img_name) # names.append(img_name) # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # convert to numpy vecs = vecs.numpy() vecs = vecs.T print("--------------------------------------------------") print(" writing feature extraction results ...") print("--------------------------------------------------") output = "gem_res_holiday_3.h5" h5f = h5py.File(output, 'w') h5f.create_dataset('dataset_1', data=vecs) h5f.create_dataset('dataset_2', data=np.string_(names)) h5f.close() print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def make_dataloader(args, config, rank=None, world_size=None): general_config = config["general"] data_config = config["dataloader"] # Manually check if there are unknown test datasets for dataset in data_config.getstruct("test_datasets"): if dataset not in test_datasets_names: raise ValueError( 'Unsupported or unknown test dataset: {}!'.format(dataset)) # Check if test datasets are available, download it if not ! download_train(args.data) download_test(args.data) # Data Loader log_debug("Creating dataloaders for dataset in %s", args.data) train_tf = TuplesTransform( shortest_size=data_config.getint("train_shortest_size"), longest_max_size=data_config.getint("train_longest_max_size"), rgb_mean=data_config.getstruct("rgb_mean"), rgb_std=data_config.getstruct("rgb_std"), random_flip=data_config.getboolean("random_flip"), random_scale=data_config.getstruct("random_scale")) # Training dataloader train_db = TuplesDataset(root_dir=args.data, name=data_config.get("training_dataset"), mode='train', neg_num=data_config.getint("neg_num"), query_size=data_config.getint("train_query_size"), pool_size=data_config.getint("train_pool_size"), transform=train_tf, batch_size=data_config.getint("train_batch_size"), num_workers=data_config.getint("num_workers")) train_sampler = TuplesDistributedARBatchSampler( train_db, batch_size=data_config.getint("train_batch_size"), num_replicas=world_size, rank=rank, drop_last=True) train_dl = data.DataLoader(train_db, batch_sampler=train_sampler, collate_fn=iss_collate_fn, pin_memory=True, num_workers=data_config.getint("num_workers"), shuffle=False) # Validation dataloader val_tf = TuplesTransform( shortest_size=data_config.getint("train_shortest_size"), longest_max_size=data_config.getint("train_longest_max_size"), rgb_mean=data_config.getstruct("rgb_mean"), rgb_std=data_config.getstruct("rgb_std"), random_flip=data_config.getboolean("random_flip"), random_scale=data_config.getstruct("random_scale")) val_db = TuplesDataset( root_dir=args.data, name=data_config.get("training_dataset"), mode='val', neg_num=data_config.getint("neg_num"), query_size=float('inf'), pool_size=float('inf'), #query_size=data_config.getint("train_query_size"), #pool_size=data_config.getint("train_pool_size"), transform=val_tf, batch_size=data_config.getint("train_batch_size"), num_workers=data_config.getint("num_workers")) val_sampler = TuplesDistributedARBatchSampler( val_db, batch_size=data_config.getint("train_batch_size"), num_replicas=world_size, rank=rank, drop_last=True) val_dl = data.DataLoader(val_db, batch_sampler=val_sampler, collate_fn=iss_collate_fn, pin_memory=True, num_workers=data_config.getint("num_workers"), shuffle=False) return train_dl, val_dl
def main(): args = parser.parse_args() # check if there are unknown datasets for dataset in args.datasets.split(','): if dataset not in datasets_names: raise ValueError( 'Unsupported or unknown dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network # pretrained networks (downloaded automatically) print(">> Loading network:\n>>>> '{}'".format(args.network)) state = load_url(PRETRAINED[args.network], model_dir=os.path.join(get_data_root(), 'networks')) # parsing net params from meta # architecture, pooling, mean, std required # the rest has default values, in case that is doesnt exist net_params = {} net_params['architecture'] = state['meta']['architecture'] net_params['pooling'] = state['meta']['pooling'] net_params['local_whitening'] = state['meta'].get('local_whitening', False) net_params['regional'] = state['meta'].get('regional', False) net_params['whitening'] = state['meta'].get('whitening', False) net_params['mean'] = state['meta']['mean'] net_params['std'] = state['meta']['std'] net_params['pretrained'] = False # network initialization net = init_network(net_params) net.load_state_dict(state['state_dict']) print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] try: bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] except: bbxs = None # for holidaysmanrot and copydays # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms) print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) top_k = 100 ranks_fnames_qs = [] for q_id in range(len(cfg["qimlist"])): ranks_q = list(ranks[:top_k, q_id]) ranks_fname_per_q = [] for img_id in ranks_q: ranks_fname_per_q.append(cfg["imlist"][img_id]) ranks_fnames_qs.append(ranks_fname_per_q) compute_map_and_print(dataset, ranks, cfg['gnd']) compute_map_and_print_top_k(dataset, ranks_fnames_qs, cfg['gnd'], cfg["imlist"]) sys.exit() with open(dataset + "_gl18_tl_resnet101_gem_w_m.pkl", "wb") as f: data = {"ranks": ranks, "db_images": images, "q_images": qimages} pickle.dump(data, f) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def main(): args = parser.parse_args() # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path if args.network_path is not None: print(">> Loading network:\n>>>> '{}'".format(args.network_path)) state = torch.load(args.network_path) net = init_network(model=state['meta']['architecture'], pooling=state['meta']['pooling'], whitening=state['meta']['whitening'], mean=state['meta']['mean'], std=state['meta']['std'], pretrained=False) net.load_state_dict(state['state_dict']) print(">>>> loaded network: ") print(net.meta_repr()) # loading offtheshelf network elif args.network_offtheshelf is not None: offtheshelf = args.network_offtheshelf.split('-') if len(offtheshelf) == 3: if offtheshelf[2] == 'whiten': offtheshelf_whiten = True else: raise (RuntimeError( "Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten" )) else: offtheshelf_whiten = False print(">> Loading off-the-shelf network:\n>>>> '{}'".format( args.network_offtheshelf)) net = init_network(model=offtheshelf[0], pooling=offtheshelf[1], whitening=offtheshelf_whiten) print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = [1] msp = 1 if args.multiscale: ms = [1, 1. / math.sqrt(2), 1. / 2] if net.meta['pooling'] == 'gem' and net.whiten is None: msp = net.pool.p.data.tolist()[0] # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # compute whitening if args.whitening is not None: start = time.time() print('>> {}: Learning whitening...'.format(args.whitening)) # loading db db_root = os.path.join(get_data_root(), 'train', args.whitening) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [ cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids'])) ] # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # learning whitening print('>> {}: Learning...'.format(args.whitening)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time() - start))) else: Lw = None datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) if dataset == 'reco': images, qimages = landmark_recognition_dataset() bbxs = [None for x in qimages] elif dataset == 'retr': images, _ = landmark_retrieval_dataset() qimages = [] bbxs = [None for x in qimages] else: # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] with open('%s_fnames.pkl' % dataset, 'wb') as f: pickle.dump([images, qimages], f) # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) vecs = vecs.numpy() print('>> saving') np.save('{}_vecs.npy'.format(dataset), vecs) if len(qimages) > 0: print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp) qvecs = qvecs.numpy() np.save('{}_qvecs.npy'.format(dataset), qvecs) if Lw is not None: # whiten the vectors vecs_lw = whitenapply(vecs, Lw['m'], Lw['P']) qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P']) # TODO print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def main(): global args, min_loss args = parser.parse_args() # manually check if there are unknown test datasets for dataset in args.test_datasets.split(','): if dataset not in test_datasets_names: raise ValueError('Unsupported or unknown test dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.training_dataset) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) if args.local_whitening: directory += "_lwhiten" if args.regional: directory += "_r" if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin) directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay) directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size) directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size) args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set cuda visible device os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # set random seeds # TODO: maybe pass as argument in future implementation? torch.manual_seed(0) torch.cuda.manual_seed_all(0) np.random.seed(0) # initialize model if args.pretrained: print(">> Using pre-trained model '{}'".format(args.arch)) else: print(">> Using model from scratch (random weights) '{}'".format(args.arch)) model_params = {} model_params['architecture'] = args.arch model_params['pooling'] = args.pool model_params['local_whitening'] = args.local_whitening model_params['regional'] = args.regional model_params['whitening'] = args.whitening # model_params['mean'] = ... # will use default # model_params['std'] = ... # will use default model_params['pretrained'] = args.pretrained model = init_network(model_params) # move network to gpu model.cuda() if model_params['pooling'] == 'netvlad': normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform = torchvision.transforms.Compose([ torchvision.transforms.Resize(size=(364, 364)), torchvision.transforms.ToTensor(), normalize, ]) image_folder = folder.ImageFolder(root="/mnt/m2/dataset", transform=transform) train_loader = torch.utils.data.DataLoader( image_folder, batch_size=64, num_workers=8, shuffle=True ) n_batches = 10 descs_list = [] i = 0 with torch.no_grad(): for x, _ in train_loader: model.eval() desc = model.compute_features(x.cuda()) max_pooled_feat_3 = torch.nn.functional.max_pool2d(desc, kernel_size=3, stride=1) max_pooled_feat_2 = torch.nn.functional.max_pool2d(desc, kernel_size=2, stride=1) reshaped_pool_3 = make_locals(max_pooled_feat_3) reshaped_pool_2 = make_locals(max_pooled_feat_2) desc = torch.cat([reshaped_pool_2, reshaped_pool_3], dim=1) # N, dim, h, w = desc.shape # desc = desc.view(N, dim, h*w).permute(0, 2, 1).reshape(N, -1, 512) desc = desc.cpu().numpy().astype('float32') descs_list.append(desc) print(">> Extracted batch {}/{} - NetVLAD initialization -".format(i+1, n_batches)) i+=1 if i == n_batches: break descs_list = np.array(descs_list).reshape(-1, 512) print(descs_list.shape) print(">> Sampling local features ") # locals = np.vstack((m[np.random.randint(len(m), size=150)] for m in descs_list)).astype('float32') locals = descs_list[np.random.randint(len(descs_list), size=len(descs_list)//3)] np.random.shuffle(locals) print(">> Locals extracted shape : {}".format(locals.shape)) n_clust = 64 locals = preprocessing.normalize(locals, axis=1) print(">> Fitting centroids with K-Means") kmeans = MiniBatchKMeans(n_clusters=n_clust).fit(locals) centroids = kmeans.cluster_centers_ print(">> Centroids shape: ", centroids.shape) model.pool.init_params(centroids.T) print(">> NetVLAD initialized") # define loss function (criterion) and optimizer if args.loss == 'contrastive': criterion = ContrastiveLoss(margin=args.loss_margin).cuda() elif args.loss == 'triplet': criterion = TripletLoss(margin=args.loss_margin).cuda() else: raise(RuntimeError("Loss {} not available!".format(args.loss))) # parameters split into features, pool, whitening # IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM parameters = [] # add feature parameters parameters.append({'params': model.features.parameters()}) # add local whitening if exists if model.lwhiten is not None: parameters.append({'params': model.lwhiten.parameters()}) # add pooling parameters (or regional whitening which is part of the pooling layer!) if not args.regional: # global, only pooling parameter p weight decay should be 0 if args.pool == 'gem': parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0}) elif args.pool == 'gemmp': parameters.append({'params': model.pool.parameters(), 'lr': args.lr*100, 'weight_decay': 0}) else: # regional, pooling parameter p weight decay should be 0, # and we want to add regional whitening if it is there if args.pool == 'gem': parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, 'weight_decay': 0}) elif args.pool == 'gemmp': parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*100, 'weight_decay': 0}) if model.pool.whiten is not None: parameters.append({'params': model.pool.whiten.parameters()}) # add final whitening if exists if model.whiten is not None: parameters.append({'params': model.whiten.parameters()}) # define optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay) # define learning rate decay schedule # TODO: maybe pass as argument in future implementation? exp_decay = math.exp(-0.01) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 if args.resume: args.resume = os.path.join(args.directory, args.resume) if os.path.isfile(args.resume): # load checkpoint weights and update model and optimizer print(">> Loading checkpoint:\n>> '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] min_loss = checkpoint['min_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) # important not to forget scheduler updating scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch']-1) else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading code normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) train_dataset = TuplesDataset( name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=args.query_size, poolsize=args.pool_size, transform=transform ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None, drop_last=True, collate_fn=collate_tuples ) if args.val: val_dataset = TuplesDataset( name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=float('Inf'), poolsize=float('Inf'), transform=transform ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=collate_tuples ) # evaluate the network before starting # this might not be necessary? test(args.test_datasets, model) for epoch in range(start_epoch, args.epochs): # set manual seeds per epoch np.random.seed(epoch) torch.manual_seed(epoch) torch.cuda.manual_seed_all(epoch) # adjust learning rate for each epoch scheduler.step() # # debug printing to check if everything ok # lr_feat = optimizer.param_groups[0]['lr'] # lr_pool = optimizer.param_groups[1]['lr'] # print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool)) # train for one epoch on train set loss = train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set if args.val: with torch.no_grad(): loss = validate(val_loader, model, criterion, epoch) # evaluate on test datasets every test_freq epochs if (epoch + 1) % args.test_freq == 0: with torch.no_grad(): test(args.test_datasets, model) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) save_checkpoint({ 'epoch': epoch + 1, 'meta': model.meta, 'state_dict': model.state_dict(), 'min_loss': min_loss, 'optimizer' : optimizer.state_dict(), }, is_best, args.directory)
def main(): global args, min_loss args = parser.parse_args() print(args) # manually check if there are unknown test datasets for dataset in args.test_datasets.split(','): if dataset not in test_datasets_names: raise ValueError( 'Unsupported or unknown test dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.training_dataset) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) if args.local_whitening: directory += "_lwhiten" if args.regional: directory += "_r" if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin) directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay) directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size) directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size) args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format( args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set cuda visible device os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # set random seeds # TODO: maybe pass as argument in future implementation? torch.manual_seed(0) torch.cuda.manual_seed_all(0) np.random.seed(0) # initialize model if args.pretrained: print(">> Using pre-trained model '{}'".format(args.arch)) else: print(">> Using model from scratch (random weights) '{}'".format( args.arch)) model_params = {} model_params['architecture'] = args.arch model_params['pooling'] = args.pool model_params['local_whitening'] = args.local_whitening model_params['regional'] = args.regional model_params['whitening'] = args.whitening model_params['mean'] = args.mean model_params['std'] = args.std model_params['pretrained'] = args.pretrained model = init_network(model_params) # move network to gpu model.cuda() print(model) summary(model, (3, 224, 224)) # define loss function (criterion) and optimizer if args.loss == 'contrastive': criterion = ContrastiveLoss(margin=args.loss_margin).cuda() elif args.loss == 'triplet': criterion = TripletLoss(margin=args.loss_margin).cuda() else: raise (RuntimeError("Loss {} not available!".format(args.loss))) # parameters split into features, pool, whitening # IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM parameters = [] # add feature parameters parameters.append({'params': model.features.parameters()}) # add local whitening if exists if model.lwhiten is not None: parameters.append({'params': model.lwhiten.parameters()}) # add pooling parameters (or regional whitening which is part of the pooling layer!) if not args.regional: # global, only pooling parameter p weight decay should be 0 if args.pool == 'gem': parameters.append({ 'params': model.pool.parameters(), 'lr': args.lr * 10, 'weight_decay': 0 }) elif args.pool == 'gemmp': parameters.append({ 'params': model.pool.parameters(), 'lr': args.lr * 100, 'weight_decay': 0 }) else: # regional, pooling parameter p weight decay should be 0, # and we want to add regional whitening if it is there if args.pool == 'gem': parameters.append({ 'params': model.pool.rpool.parameters(), 'lr': args.lr * 10, 'weight_decay': 0 }) elif args.pool == 'gemmp': parameters.append({ 'params': model.pool.rpool.parameters(), 'lr': args.lr * 100, 'weight_decay': 0 }) if model.pool.whiten is not None: parameters.append({'params': model.pool.whiten.parameters()}) # add final whitening if exists if model.whiten is not None: parameters.append({'params': model.whiten.parameters()}) # define optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay) # define learning rate decay schedule # TODO: maybe pass as argument in future implementation? exp_decay = math.exp(-0.01) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 if args.resume: args.resume = os.path.join(args.directory, args.resume) if os.path.isfile(args.resume): # load checkpoint weights and update model and optimizer print(">> Loading checkpoint:\n>> '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] min_loss = checkpoint['min_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) # important not to forget scheduler updating scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1) else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading code normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) train_dataset = TuplesDataset(name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=args.query_size, poolsize=args.pool_size, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None, drop_last=True, collate_fn=collate_tuples) if args.val: val_dataset = TuplesDataset(name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=float('Inf'), poolsize=float('Inf'), transform=transform) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=collate_tuples) # uncomment to evaluate the network before starting # test(args.test_datasets, model) for epoch in range(start_epoch, args.epochs): # set manual seeds per epoch np.random.seed(epoch) torch.manual_seed(epoch) torch.cuda.manual_seed_all(epoch) # adjust learning rate for each epoch scheduler.step() # # debug printing to check if everything ok # lr_feat = optimizer.param_groups[0]['lr'] # lr_pool = optimizer.param_groups[1]['lr'] # print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool)) # train for one epoch on train set # loss = train(train_loader, model, criterion, optimizer, epoch) loss = 999. # evaluate on validation set if args.val: with torch.no_grad(): loss = validate(val_loader, model, criterion, epoch) # evaluate on test datasets every test_freq epochs if (epoch + 1) % args.test_freq == 0: with torch.no_grad(): test(args.test_datasets, model) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) save_checkpoint( { 'epoch': epoch + 1, 'meta': model.meta, 'state_dict': model.state_dict(), 'min_loss': min_loss, 'optimizer': optimizer.state_dict(), }, is_best, args.directory)
def main(): global args, min_loss args = parser.parse_args() # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.training_dataset) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin) directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay) directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size) directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size) args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format( args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set cuda visible device os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # set random seeds (maybe pass as argument) torch.manual_seed(0) torch.cuda.manual_seed_all(0) np.random.seed(0) # create model if args.pretrained: print(">> Using pre-trained model '{}'".format(args.arch)) model = init_network(model=args.arch, pooling=args.pool, whitening=args.whitening) else: print(">> Using model from scratch (random weights) '{}'".format( args.arch)) model = init_network(model=args.arch, pooling=args.pool, whitening=args.whitening, pretrained=False) # move network to gpu model.cuda() # define loss function (criterion) and optimizer if args.loss == 'contrastive': criterion = ContrastiveLoss(margin=args.loss_margin).cuda() else: raise (RuntimeError("Loss {} not available!".format(args.loss))) # parameters split into features and pool (no weight decay for pooling layer) parameters = [{ 'params': model.features.parameters() }, { 'params': model.pool.parameters(), 'lr': args.lr * 10, 'weight_decay': 0 }] if model.whiten is not None: parameters.append({'params': model.whiten.parameters()}) # define optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay) # define learning rate decay schedule exp_decay = math.exp(-0.01) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 if args.resume: args.resume = os.path.join(args.directory, args.resume) if os.path.isfile(args.resume): print(">> Loading checkpoint:\n>> '{}'".format(args.resume)) checkpoint = torch.load(args.resume) start_epoch = checkpoint['epoch'] min_loss = checkpoint['min_loss'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1) else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading code normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) transform = transforms.Compose([ transforms.ToTensor(), normalize, ]) train_dataset = TuplesDataset(name=args.training_dataset, mode='train', imsize=args.image_size, nnum=args.neg_num, qsize=args.query_size, poolsize=args.pool_size, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, sampler=None, drop_last=True, collate_fn=collate_tuples) if args.val: val_dataset = TuplesDataset(name=args.training_dataset, mode='val', imsize=args.image_size, nnum=args.neg_num, qsize=float('Inf'), poolsize=float('Inf'), transform=transform) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=collate_tuples) # evaluate the network before starting test(args.test_datasets, model) for epoch in range(start_epoch, args.epochs): # set manual seeds per epoch np.random.seed(epoch) torch.manual_seed(epoch) torch.cuda.manual_seed_all(epoch) # adjust learning rate for each epoch scheduler.step() # lr_feat = optimizer.param_groups[0]['lr'] # lr_pool = optimizer.param_groups[1]['lr'] # print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool)) # train for one epoch on train set loss = train(train_loader, model, criterion, optimizer, epoch) # evaluate on validation set if args.val: loss = validate(val_loader, model, criterion, epoch) # evaluate on test datasets test(args.test_datasets, model) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) save_checkpoint( { 'epoch': epoch + 1, 'meta': model.meta, 'state_dict': model.state_dict(), 'min_loss': min_loss, 'optimizer': optimizer.state_dict(), }, is_best, args.directory)
def main(): global args, min_loss args = parser.parse_args() # set cuda visible device os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # set random seeds (maybe pass as argument) torch.manual_seed(1234) torch.cuda.manual_seed_all(1234) np.random.seed(1234) random.seed(1234) torch.backends.cudnn.deterministic = True state = torch.load(args.network_path) model = init_network( model=state["meta"]["architecture"], pooling=state["meta"]["pooling"], whitening=state["meta"]["whitening"], mean=state["meta"]["mean"], std=state["meta"]["std"], pretrained=False, ) model.load_state_dict(state["state_dict"]) model.meta["Lw"] = state["meta"]["Lw"] model.cuda() # whitening Lw = model.meta["Lw"]["retrieval-SfM-120k"]["ss"] Lw_m = torch.from_numpy(Lw["m"]).cuda().float() Lw_p = torch.from_numpy(Lw["P"]).cuda().float() # Data loading code normalize = transforms.Normalize(mean=model.meta["mean"], std=model.meta["std"]) transform = transforms.Compose([transforms.ToTensor(), normalize]) query_dataset = MyTripletDataset( # imsize=args.image_size, imsize=(362, 362), random=False, transform=transform, norm=(0, 0), filename="base/" + args.network_path.replace("/", "_") + "_triplet", ) # val_dataset.test_cluster(model) # return query_dataset.create_epoch_tuples(model) query_loader = torch.utils.data.DataLoader( query_dataset, batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True, worker_init_fn=lambda _: random.seed(1234), ) base_dataset = ImagesFromList( root="", images=query_dataset.images, # imsize=query_dataset.imsize, imsize=(362, 362), transform=query_dataset.transform, random=False, ) base_loader = torch.utils.data.DataLoader(base_dataset, batch_size=1, shuffle=False) # test(["oxford5k"], model) extract(query_loader, base_loader, model, Lw_m, Lw_p)
def main(): args = parser.parse_args() # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path if args.network_path is not None: print(">> Loading network:\n>>>> '{}'".format(args.network_path)) if args.network_path in PRETRAINED: # pretrained networks (downloaded automatically) state = load_url(PRETRAINED[args.network_path], model_dir=os.path.join(get_data_root(), 'networks')) else: state = torch.load(args.network_path) net = init_network(model=state['meta']['architecture'], pooling=state['meta']['pooling'], whitening=state['meta']['whitening'], mean=state['meta']['mean'], std=state['meta']['std'], pretrained=False) net.load_state_dict(state['state_dict']) # if whitening is precomputed if 'Lw' in state['meta']: net.meta['Lw'] = state['meta']['Lw'] print(">>>> loaded network: ") print(net.meta_repr()) # loading offtheshelf network elif args.network_offtheshelf is not None: offtheshelf = args.network_offtheshelf.split('-') if len(offtheshelf) == 3: if offtheshelf[2] == 'whiten': offtheshelf_whiten = True else: raise (RuntimeError( "Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten" )) else: offtheshelf_whiten = False print(">> Loading off-the-shelf network:\n>>>> '{}'".format( args.network_offtheshelf)) net = init_network(model=offtheshelf[0], pooling=offtheshelf[1], whitening=offtheshelf_whiten) print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = [1] msp = 1 if args.multiscale: ms = [1, 1. / math.sqrt(2), 1. / 2] if net.meta['pooling'] == 'gem' and net.whiten is None: msp = net.pool.p.data.tolist()[0] # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize(mean=net.meta['mean'], std=net.meta['std']) transform = transforms.Compose([transforms.ToTensor(), normalize]) # compute whitening if args.whitening is not None: start = time.time() if 'Lw' in net.meta and args.whitening in net.meta['Lw']: print('>> {}: Whitening is precomputed, loading it...'.format( args.whitening)) if args.multiscale: Lw = net.meta['Lw'][args.whitening]['ms'] else: Lw = net.meta['Lw'][args.whitening]['ss'] else: print('>> {}: Learning whitening...'.format(args.whitening)) # loading db db_root = os.path.join(get_data_root(), 'train', args.whitening) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [ cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids'])) ] # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # learning whitening print('>> {}: Learning...'.format(args.whitening)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time() - start))) else: Lw = None # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])] bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp) print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset, ranks, cfg['gnd']) if Lw is not None: # whiten the vectors vecs_lw = whitenapply(vecs, Lw['m'], Lw['P']) qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P']) # search, rank, and print scores = np.dot(vecs_lw.T, qvecs_lw) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd']) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
def main(): args = parser.parse_args() # check if there are unknown datasets for dataset in args.datasets.split(','): if dataset not in datasets_names: raise ValueError('Unsupported or unknown dataset: {}!'.format(dataset)) # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path if args.network_path is not None: print(">> Loading network:\n>>>> '{}'".format(args.network_path)) if args.network_path in PRETRAINED: # pretrained networks (downloaded automatically) state = load_url(PRETRAINED[args.network_path], model_dir=os.path.join(get_data_root(), 'networks')) else: # fine-tuned network from path state = torch.load(args.network_path) # parsing net params from meta # architecture, pooling, mean, std required # the rest has default values, in case that is doesnt exist net_params = {} net_params['architecture'] = state['meta']['architecture'] net_params['pooling'] = state['meta']['pooling'] net_params['local_whitening'] = state['meta'].get('local_whitening', False) net_params['regional'] = state['meta'].get('regional', False) net_params['whitening'] = state['meta'].get('whitening', False) net_params['mean'] = state['meta']['mean'] net_params['std'] = state['meta']['std'] net_params['pretrained'] = False # load network net = init_network(net_params) net.load_state_dict(state['state_dict']) # if whitening is precomputed if 'Lw' in state['meta']: net.meta['Lw'] = state['meta']['Lw'] print(">>>> loaded network: ") print(net.meta_repr()) # loading offtheshelf network elif args.network_offtheshelf is not None: # parse off-the-shelf parameters offtheshelf = args.network_offtheshelf.split('-') net_params = {} net_params['architecture'] = offtheshelf[0] net_params['pooling'] = offtheshelf[1] net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:] net_params['regional'] = 'reg' in offtheshelf[2:] net_params['whitening'] = 'whiten' in offtheshelf[2:] net_params['pretrained'] = True # load off-the-shelf network print(">> Loading off-the-shelf network:\n>>>> '{}'".format(args.network_offtheshelf)) net = init_network(net_params) print(">>>> loaded network: ") print(net.meta_repr()) # setting up the multi-scale parameters ms = list(eval(args.multiscale)) if len(ms)>1 and net.meta['pooling'] == 'gem' and not net.meta['regional'] and not net.meta['whitening']: msp = net.pool.p.item() print(">> Set-up multiscale:") print(">>>> ms: {}".format(ms)) print(">>>> msp: {}".format(msp)) else: msp = 1 # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize( mean=net.meta['mean'], std=net.meta['std'] ) transform = transforms.Compose([ transforms.ToTensor(), normalize ]) # compute whitening if args.whitening is not None: start = time.time() if 'Lw' in net.meta and args.whitening in net.meta['Lw']: print('>> {}: Whitening is precomputed, loading it...'.format(args.whitening)) if len(ms)>1: Lw = net.meta['Lw'][args.whitening]['ms'] else: Lw = net.meta['Lw'][args.whitening]['ss'] else: # if we evaluate networks from path we should save/load whitening # not to compute it every time if args.network_path is not None: whiten_fn = args.network_path + '_{}_whiten'.format(args.whitening) if len(ms) > 1: whiten_fn += '_ms' whiten_fn += '.pth' else: whiten_fn = None if whiten_fn is not None and os.path.isfile(whiten_fn): print('>> {}: Whitening is precomputed, loading it...'.format(args.whitening)) Lw = torch.load(whiten_fn) else: print('>> {}: Learning whitening...'.format(args.whitening)) # loading db db_root = os.path.join(get_data_root(), 'train', args.whitening) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))] # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # learning whitening print('>> {}: Learning...'.format(args.whitening)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} # saving whitening if whiten_fn exists if whiten_fn is not None: print('>> {}: Saving to {}...'.format(args.whitening, whiten_fn)) torch.save(Lw, whiten_fn) print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time()-start))) else: Lw = None # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) images = [cfg['im_fname'](cfg,i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg,i) for i in range(cfg['nq'])] try: bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] except: bbxs = None # for holidaysmanrot and copydays # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp) print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) top_k = -1 ranks_fnames_qs = [] for q_id in range(len(cfg["qimlist"])): ranks_q = list(ranks[:, q_id]) if top_k == -1 else list(ranks[:top_k, q_id]) ranks_fname_per_q = [] for img_id in ranks_q: ranks_fname_per_q.append(cfg["imlist"][img_id]) ranks_fnames_qs.append(ranks_fname_per_q) compute_map_and_print(dataset, ranks, cfg['gnd']) compute_map_and_print_top_k(dataset, ranks_fnames_qs, cfg['gnd'], cfg["imlist"]) with open(dataset + "_topAll_resnet101_sfM120k_gem_ms.pkl", "wb") as f: pickle.dump(ranks_fnames_qs, f) sys.exit() if Lw is not None: # whiten the vectors vecs_lw = whitenapply(vecs, Lw['m'], Lw['P']) qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P']) # search, rank, and print scores = np.dot(vecs_lw.T, qvecs_lw) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd']) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time()-start)))
def main(): global args, min_loss args = parser.parse_args() print(args) # set cuda visible device os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # create export dir if it doesnt exist directory = "{}".format(args.target.replace("/", "_")) directory += "_{}".format(args.arch) directory += "_{}".format(args.pool) if args.whitening: directory += "_whiten" if not args.pretrained: directory += "_notpretrained" # directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size) directory += "_pretrained" if args.is_pretrained else "" directory += "_random" if args.is_random else "" directory += args.notion target_net = args.target[6:].replace("_", "/") print(target_net) state = torch.load(target_net) lw = state["meta"]["Lw"]["retrieval-SfM-120k"]["ss"] args.directory = os.path.join(args.directory, directory) print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) # set random seeds (maybe pass as argument) torch.manual_seed(0) torch.cuda.manual_seed_all(0) np.random.seed(0) # create model print(">> Using pre-trained model '{}'".format(args.arch)) model = init_network( model=args.arch, pooling=args.pool, whitening=args.whitening, pretrained=not args.is_random, ) model.cuda() target_model = init_network( model=args.arch, pooling=args.pool, whitening=args.whitening, pretrained=not args.is_random, ) target_model.load_state_dict(state["state_dict"]) lw_m = lw["m"].copy() lw_p = lw["P"].copy() target_model.lw_m = nn.Parameter(torch.from_numpy(lw_m).float()) target_model.lw_p = nn.Parameter(torch.from_numpy(lw_p).float()) target_model.cuda() lw_m = lw["m"].copy() lw_p = lw["P"].copy() model.lw_m = nn.Parameter(torch.from_numpy(lw_m).float()) model.lw_p = nn.Parameter(torch.from_numpy(lw_p).float()) whiten_layer = Whiten_layer(lw["P"].shape[1], lw["P"].shape[0]) model.white_layer = whiten_layer model.cuda() # parameters split into features and pool (no weight decay for pooling layer) parameters = [ {"params": model.features.parameters()}, {"params": model.pool.parameters(), "lr": args.lr * 10, "weight_decay": 0}, {"params": model.white_layer.parameters(), "lr": 1e-2, "weight_decay": 5e-1}, ] if model.whiten is not None: parameters.append({"params": model.whiten.parameters()}) # define optimizer if args.optimizer == "sgd": optimizer = torch.optim.SGD( parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay ) elif args.optimizer == "adam": optimizer = torch.optim.Adam( parameters, args.lr, weight_decay=args.weight_decay ) # define learning rate decay schedule exp_decay = math.exp(-0.01) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay) # optionally resume from a checkpoint start_epoch = 0 # Data loading code normalize = transforms.Normalize(mean=model.meta["mean"], std=model.meta["std"]) transform = transforms.Compose([transforms.ToTensor(), normalize]) val_dataset = Distillation_dataset( imsize=(args.image_size, args.image_size), nnum=1, qsize=float("Inf"), poolsize=float("Inf"), transform=transform, filename=args.target, q_percent=args.q_percent, ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=10, shuffle=False, num_workers=args.workers, pin_memory=False, collate_fn=collate_tuples, ) min_epoch = -1 for epoch in range(start_epoch, args.epochs): if args.is_pretrained or args.is_random: break # set manual seeds per epoch np.random.seed(epoch) torch.manual_seed(epoch) torch.cuda.manual_seed_all(epoch) # adjust learning rate for each epoch scheduler.step() loss = train(val_loader, model, optimizer, epoch, target_model) print(loss) # evaluate on test datasets if (epoch + 1) % 1 == 0: with torch.no_grad(): test(args.test_datasets, model, lw) # remember best loss and save checkpoint is_best = loss < min_loss min_loss = min(loss, min_loss) save_checkpoint( { "epoch": epoch + 1, "meta": model.meta, "state_dict": model.state_dict(), "min_loss": min_loss, "optimizer": optimizer.state_dict(), }, is_best, args.directory, ) if is_best: min_epoch = epoch # if epoch - min_epoch > 5: # # break # if val_dataset.phase == 1: # print(bcolors.str(">>> phase 2", bcolors.OKGREEN)) # val_dataset.phase = 2 # min_epoch = epoch # for group in optimizer.param_groups: # group["lr"] /= 10 # else: # break if args.is_pretrained or args.is_random: save_checkpoint( { "epoch": 0 + 1, "meta": model.meta, "state_dict": model.state_dict(), "min_loss": min_loss, "optimizer": optimizer.state_dict(), }, True, args.directory, )
def main(): args = parser.parse_args() # check if test dataset are downloaded # and download if they are not download_train(get_data_root()) download_test(get_data_root()) # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path if args.network_path is not None: net = load_network(args.network_path) # loading offtheshelf network elif args.network_offtheshelf is not None: net = load_offtheshelf(args.network_offtheshelf) # setting up the multi-scale parameters ms = [1] msp = 1 if args.multiscale: ms = [1, 1./math.sqrt(2), 1./2] if net.meta['pooling'] == 'gem' and net.whiten is None: msp = net.pool.p.data.tolist()[0] # moving network to gpu and eval mode net.cuda() net.eval() # set up the transform normalize = transforms.Normalize( mean=net.meta['mean'], std=net.meta['std'] ) transform = transforms.Compose([ transforms.ToTensor(), normalize ]) # compute whitening if args.whitening is not None: start = time.time() if 'Lw' in net.meta and args.whitening in net.meta['Lw']: print('>> {}: Whitening is precomputed, loading it...'.format(args.whitening)) if args.multiscale: Lw = net.meta['Lw'][args.whitening]['ms'] else: Lw = net.meta['Lw'][args.whitening]['ss'] else: print('>> {}: Learning whitening...'.format(args.whitening)) if args.whitening == "scores": # special logic for scores database from score_retrieval.exports import ( db, train_images as images, ) else: # loading db db_root = os.path.join(get_data_root(), 'train', args.test_whiten) ims_root = os.path.join(db_root, 'ims') db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.test_whiten)) with open(db_fn, 'rb') as f: db = pickle.load(f) images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))] # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # learning whitening print('>> {}: Learning...'.format(args.whitening)) wvecs = wvecs.numpy() m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs']) Lw = {'m': m, 'P': P} print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time()-start))) else: Lw = None # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) if dataset == "scores": # Special added logic to handle loading our score dataset from score_retrieval.exports import ( images, qimages, gnd, ) print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, ms=ms, msp=msp) else: # extract ground truth cfg = configdataset(dataset, os.path.join(get_data_root(), 'test')) gnd = cfg['gnd'] # prepare config structure for the test dataset images = [cfg['im_fname'](cfg,i) for i in range(cfg['n'])] qimages = [cfg['qim_fname'](cfg,i) for i in range(cfg['nq'])] bbxs = [tuple(gnd[i]['bbx']) for i in range(cfg['nq'])] # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp) # validation print(">> {}: gnd stats: {}, {}, {}".format( dataset, len(gnd), [len(x["ok"]) for x in gnd[10:]], [len(x["junk"]) for x in gnd[10:]], )) print(">> {}: image stats: {}, {}".format(dataset, len(images), len(qimages))) assert len(gnd) == len(qimages), (len(gnd), len(qimages)) print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() print(">> {}: qvecs.shape: {}".format(dataset, qvecs.shape)) # search, rank, and print scores = np.dot(vecs.T, qvecs) ranks = np.argsort(-scores, axis=0) print(">> {}: ranks (shape {}) head: {}".format(dataset, ranks.shape, ranks[10:,10:])) print(">> {}: gnd head: {}".format(dataset, gnd[5:])) # Compute and print metrics compute_acc(ranks, gnd, dataset) compute_mrr(ranks, gnd, dataset) compute_map_and_print(dataset, ranks, gnd) if Lw is not None: # whiten the vectors vecs_lw = whitenapply(vecs, Lw['m'], Lw['P']) qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P']) # search, rank, and print scores = np.dot(vecs_lw.T, qvecs_lw) ranks = np.argsort(-scores, axis=0) compute_acc(ranks, gnd, dataset + " + whiten") compute_mrr(ranks, gnd, dataset + " + whiten") compute_map_and_print(dataset + " + whiten", ranks, gnd) print('>> {}: elapsed time: {}'.format(dataset, htime(time.time()-start)))