def loadModelToCuda(eConfig): kargs = { 'model': eConfig['arch'], 'pooling': eConfig['pooling'], 'whitening': False } #eConfig['whitening'] } if (eConfig['pretrained'] == 'imageNet'): model = init_network(**kargs, pretrained=True) else: model = init_network(**kargs, pretrained=False) return model.cuda()
def constructfeature(self, hash_size, input_dim, num_hashtables): multiscale = '[1]' print(">> Loading network:\n>>>> '{}'".format(self.network)) # state = load_url(PRETRAINED[args.network], model_dir=os.path.join(get_data_root(), 'networks')) state = torch.load(self.network) # 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(multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode if torch.cuda.is_available(): 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 ]) # extract database and query vectors print('>> database images...') images = ImageProcess(self.img_dir).process() vecs, img_paths = extract_vectors(net, images, 1024, transform, ms=ms) feature_dict = dict(zip(img_paths, list(vecs.detach().cpu().numpy().T))) # index lsh = LSHash(hash_size=int(hash_size), input_dim=int(input_dim), num_hashtables=int(num_hashtables)) for img_path, vec in feature_dict.items(): lsh.index(vec.flatten(), extra_data=img_path) # ## 保存索引模型 # with open(self.feature_path, "wb") as f: # pickle.dump(feature_dict, f) # with open(self.index_path, "wb") as f: # pickle.dump(lsh, f) print("extract feature is done") return feature_dict, lsh
def load_network(model_dir: str, device: torch.device) -> ImageRetrievalNet: state = load_url(PRETRAINED['rSfM120k-tl-resnet101-gem-w'], model_dir=model_dir) net = init_network({ 'architecture': state['meta']['architecture'], 'pooling': state['meta']['pooling'], 'whitening': state['meta'].get('whitening', False) }) net.load_state_dict(state['state_dict']) net.to(device) net.eval() log.info(f"Loaded network: {net.meta_repr()}") return net
def test_feature(self): multiscale = '[1]' print(">> Loading network:\n>>>> '{}'".format(self.network)) # state = load_url(PRETRAINED[args.network], model_dir=os.path.join(get_data_root(), 'networks')) state = torch.load(self.network) # 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(multiscale)) print(">>>> Evaluating scales: {}".format(ms)) # moving network to gpu and eval mode if torch.cuda.is_available(): 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 ]) # extract database and query vectors print('>> database images...') images = ImageProcess(self.img_dir).process() vecs, img_paths = extract_vectors(net, images, 1024, transform, ms=ms) feature_dict = dict(zip(img_paths, list(vecs.detach().cpu().numpy().T))) return feature_dict
def load_offtheshelf(network_name): """Load off the shelf network.""" offtheshelf = network_name.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(network_name)) net = init_network(model=offtheshelf[0], pooling=offtheshelf[1], whitening=offtheshelf_whiten) print(">>>> loaded network: ") print(net.meta_repr()) return net
def load_network(network_path): """Load network from a path.""" print(">> Loading network:\n>>>> '{}'".format(network_path)) if network_path in PRETRAINED: # pretrained networks (downloaded automatically) state = load_url(PRETRAINED[network_path], model_dir=os.path.join(get_data_root(), 'networks')) else: state = torch.load(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()) return net
def build_retrievalnet_from_options(opt, is_cuda=False): logger = logging.getLogger("detector-retrieval.build_retrievalnet") logger.info("Building the retrieval model...") logger.info("Loading weights from {}".format(opt.retrieval_network_path)) state = torch.load(opt.retrieval_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"] logger.info("Loaded network: ") if "epoch" in state: logger.info("Model after {} epochs".format(state["epoch"])) logger.info(net.meta_repr()) if is_cuda: net.cuda() else: net.cpu() net.eval() return net
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, max_meter args = parser.parse_args() # 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{}_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 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() # 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, 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 parameters.append({ 'params': model.pool.parameters(), 'lr': args.lr * 10, 'weight_decay': 0 }) else: # regional, pooling parameter p weight decay should be 0, # and we want to add regional whitening if it is there parameters.append({ 'params': model.pool.rpool.parameters(), 'lr': args.lr * 10, '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'] max_meter = checkpoint['max_meter'] 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) 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) # 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 % args.test_freq == 0: with torch.no_grad(): cur_meter = test(args.test_datasets, model, image_size=args.image_size) # remember best loss and save checkpoint is_best = cur_meter > max_meter max_meter = max(cur_meter, max_meter) save_checkpoint( { 'epoch': epoch, 'meta': model.meta, 'state_dict': model.state_dict(), 'max_meter': max_meter, 'optimizer': optimizer.state_dict(), }, is_best, args.directory, save_regular=args.save_freq and epoch % args.save_freq == 0 or epoch == 0) # train for one epoch on train set loss = train(train_loader, model, criterion, optimizer, epoch) # adjust learning rate for each epoch scheduler.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 net_params['multi_layer_cat'] = state['meta']['multi_layer_cat'] # 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 print(">> image size: {}".format(args.image_size)) 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 print(">> Set-up multiscale:") print(">>>> ms: {}".format(ms)) print(">>>> msp: {}".format(msp)) # 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'])] # bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] print('>> not use bbxs...') bbxs = None # key_url_list = ParseData(os.path.join(get_data_root(), 'index.csv')) # index_image_path = os.path.join(get_data_root(), 'resize_index_image') # images = [os.path.join(index_image_path, key_url_list[i][0]) for i in range(len(key_url_list))] # key_url_list = ParseData(os.path.join(get_data_root(), 'test.csv')) # test_image_path = os.path.join(get_data_root(), 'resize_test_image') # qimages = [os.path.join(test_image_path, key_url_list[i][0]) for i in range(len(key_url_list))] # # bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])] # csvfile = open(os.path.join(get_data_root(), 'index_clear.csv'), 'r') # csvreader = csv.reader(csvfile) # images = [line[:1][0] for line in csvreader] # # csvfile = open(os.path.join(get_data_root(), 'test_clear.csv'), 'r') # csvreader = csv.reader(csvfile) # qimages = [line[:1][0] for line in csvreader] # bbxs = None # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp) # vecs = torch.randn(2048, 5063) # vecs = torch.randn(2048, 4993) # hxq modified # bbxs = None # print('>> set no bbxs...') 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) # hxq modified, test add features map for retrieval # vecs = [vecs[i].numpy() for i in range(len(vecs))] # qvecs_temp = np.zeros((qvecs[0].shape[0], len(qvecs))) # for i in range(len(qvecs)): # qvecs_temp[:, i] = qvecs[i][:, 0].numpy() # qvecs = qvecs_temp # # scores = np.zeros((len(vecs), qvecs.shape[-1])) # for i in range(len(vecs)): # scores[i, :] = np.amax(np.dot(vecs[i].T, qvecs), 0) ranks = np.argsort(-scores, axis=0) mismatched_info = compute_map_and_print(dataset, ranks, cfg['gnd'], kappas=[1, 5, 10, 100]) # hxq added show_false_img = False if show_false_img == True: print('>> Save mismatched image tuple...') for info in mismatched_info: mismatched_img_show_save(info, qimages, images, args, bbxs=bbxs) 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) mismatched_info = compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd']) # hxq added # show_false_img = False if show_false_img == True: print('>> Save mismatched image tuple...') for info in mismatched_info: mismatched_img_show_save(info, qimages, images, args, bbxs=bbxs) 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: 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() # setting up the visible GPU os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id # loading network from path result_dir = 'retreival_results' if args.network_path is not None: result_dir = os.path.join(result_dir, args.network_path) 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: result_dir = os.path.join(result_dir, args.network_offtheshelf) 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: # Save whitening TODO 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 data_root = args.data_root datasets = datasets_names[args.dataset] result_dict = {} for dataset in datasets: start = time.time() result_dict[dataset] = {} print('>> {}: Extracting...'.format(dataset)) # prepare config structure for the test dataset images = get_imlist(data_root, dataset, args.train_txt) qimages = get_imlist(data_root, dataset, args.query_txt) # extract database and query vectors print('>> {}: database images...'.format(dataset)) vecs = extract_vectors(net, images, args.image_size, transform, root=os.path.join(data_root, dataset), ms=ms, msp=msp) print('>> {}: query images...'.format(dataset)) qvecs = extract_vectors(net, qimages, args.image_size, transform, root=os.path.join(data_root, dataset), ms=ms, msp=msp) print('>> {}: Evaluating...'.format(dataset)) # convert to numpy vecs = vecs.numpy() qvecs = qvecs.numpy() scores, ranks = cal_ranks(vecs, vecs, Lw) result_dict[dataset]['train'] = {'scores': scores, 'ranks': ranks} scores, ranks = cal_ranks(vecs, qvecs, Lw) result_dict[dataset]['test'] = {'scores': scores, 'ranks': ranks} print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start))) # Save retrieval results if not os.path.exists(result_dir): os.makedirs(result_dir) result_file = os.path.join(result_dir, args.outfile) np.save(result_file, result_dict) print('Save retrieval results to {}'.format(result_file))
def main(): args = parser.parse_args() # check if there are unknown datasets for scene in args.scenes.split(','): if scene not in datasets_names: raise ValueError('Unsupported or unknown scene: {}!'.format(scene)) # 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 scenes = args.scenes.split(',') for scene in scenes: start = time.time() print('>> {}: Extracting...'.format(scene)) img_path = osp.join(args.data_path, scene, "images") images = [ osp.join(img_path, fname) for fname in os.listdir(img_path) if fname[-3:].lower() in ['jpg', 'png'] ] # extract vectors vecs = extract_vectors(net, images, args.image_size, transform, ms=ms) print('>> {}: Evaluating...'.format(scene)) # convert to numpy vecs = vecs.numpy() # search, rank, and print scores = np.dot(vecs.T, vecs) ranks = np.argsort(-scores, axis=0) images = [img.split('/')[-1] for img in images] for top_k in list(eval(args.top_n)): pairs = [] for q_id in range(len(images)): img_q = images[q_id] pairs_per_q = [ " ".join([img_q, images[db_id]]) for db_id in list(ranks[1:top_k + 1, q_id]) ] pairs += pairs_per_q with open( osp.join(args.data_path, scene, "image_pairs_" + str(top_k) + ".txt"), "w") as f: for pair in pairs: f.write(pair + "\n") print('>> {}: elapsed time: {}'.format(scene, htime(time.time() - start)))
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(): arg_parser = argparse.ArgumentParser() arg_parser.add_argument('train_images_path') arg_parser.add_argument('test_images_path') arg_parser.add_argument('predictions_path') args = arg_parser.parse_args() imsize=480 train_path=args.train_images_path test_path=args.test_images_path outfile=args.predictions_path ##read data## data_list=os.listdir(ipath) train_images = get_imlist(train_path) test_images = get_imlist(test_path) ##RAMAC## RAMAC = extract_feature(train_images, 'resnet101', imsize) RAMAC_test = extract_feature(test_images, 'resnet101', imsize) ##UEL## normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform_train = transforms.Compose([ transforms.ToPILImage(), transforms.RandomCrop(size=224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), normalize, ]) transform_test = transforms.Compose([ transforms.ToPILImage(), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]) net = resnet101(pretrained=True,low_dim=128) model_path = './model/UEL.t'#After training UEL net.load_state_dict(torch.load()) imset = DataLoader(path = train_path, transform=transform_test) train_loader = torch.utils.data.DataLoader(imset, batch_size=32, shuffle=False, num_workers=0) UEL = obtainf(net, train_loader) imset = DataLoader(path = test_path, transform=transform_test) test_loader = torch.utils.data.DataLoader(imset, batch_size=32, shuffle=False, num_workers=0) UEL_test = obtainf(net, test_loader) ##GEM## image_size=1024 multiscale='[1, 2**(1/2), 1/2**(1/2)]' state = torch.load('./model/retrievalSfM120k-vgg16-gem-b4dcdc6.pth') 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'] ms = list(eval(multiscale)) msp = net.pool.p.item() 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 ]) GEM = extract_vectors(net,train_images , 480, transform, ms=ms, msp=msp).numpy().T GEM_test = extract_vectors(net,test_images , 480, transform, ms=ms, msp=msp).numpy().T ##Retrieval## feats=np.concatenate((RAMAC,UEL,GEM),axis=1).astype('float32') query_feat=np.concatenate((RAMAC_test, UEL_test,GEM_test),axis=1).astype('float32') ##diffusion## kq, kd = 7, 50 gamma=80 diffusion = Diffusion(feats, '/') offline = diffusion.get_offline_results(1024, kd) print('[search] 1) k-NN search') sims, ids = diffusion.knn.search(query_feat, kq) sims = sims ** gamma qr_num = ids.shape[0] print('[search] 2) linear combination') all_scores = np.empty((qr_num, 7), dtype=np.float32) all_ranks = np.empty((qr_num, 7), dtype=np.int) for i in range(qr_num): scores = sims[i] @ offline[ids[i]] parts = np.argpartition(-scores, 7)[:7] ranks = np.argsort(-scores[parts]) all_scores[i] = scores[parts][ranks] all_ranks[i] = parts[ranks] I = all_ranks ##output## out=pd.DataFrame(list(map(lambda x: x.split('/')[-1].split('.jpg')[0],timages))) out['1']=pd.DataFrame(I)[0].map(lambda x:data_list[x].split('.')[0] ) out['2']=pd.DataFrame(I)[1].map(lambda x:data_list[x].split('.')[0] ) out['3']=pd.DataFrame(I)[2].map(lambda x:data_list[x].split('.')[0] ) out['4']=pd.DataFrame(I)[3].map(lambda x:data_list[x].split('.')[0] ) out['5']=pd.DataFrame(I)[4].map(lambda x:data_list[x].split('.')[0] ) out['6']=pd.DataFrame(I)[5].map(lambda x:data_list[x].split('.')[0] ) out['7']=pd.DataFrame(I)[6].map(lambda x:data_list[x].split('.')[0] ) out.to_csv(outfile,index=None,header=None)
N, C, H, W = 10, 3, 256, 256 D_out = 2048 x = torch.randn(N, C, H, W) y = torch.randn(D_out, N) model_params = {} model_params['architecture'] = 'resnet101' model_params['pooling'] = 'gem' model_params['local_whitening'] = False model_params['regional'] = False model_params['whitening'] = False # model_params['mean'] = ... # will use default # model_params['std'] = ... # will use default model_params['pretrained'] = True model = init_network(model_params) # model = torch.nn.Sequential( # DataParallelModel(D_in, H), # DataParallelModel(H, H), # torch.nn.Linear(H, D_out) # ) model.cuda() criterion = ContrastiveLoss(margin=0.85).cuda() optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) # # Data loading code # normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std']) # transform = transforms.Compose([
def main(): global args, min_loss args = parser.parse_args() # create export dir if it doesnt exist conftxt = 'Model {}-{}'.format(args.arch, args.pool) if args.whitening: conftxt += '-whiten' if args.pretrained in PRETRAINED: conftxt += " Pretrained: {} ".format(args.pretrained) elif args.pretrained == 'imagenet': conftxt += " Pretrained: imagenet" else: conftxt += " No pretrain " conftxt += " Loss: {} margin: {:.2f}".format(args.loss, args.loss_margin) conftxt += " Opt: {} lr{:.1e} wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay) conftxt += " snn{} dnn{}".format(args.sneg_num, args.dneg_num) conftxt += " qsize{}".format(args.query_size) conftxt += " bsize{} imsize{}".format(args.batch_size, args.image_size) print(">> Creating directory if it does not exist:\n>> '{}'".format( args.directory)) if not os.path.exists(args.directory): os.makedirs(args.directory) log = open(os.path.join(args.directory, 'log.txt'), 'w') lprint('>>>>Training configuration:\n {}'.format(conftxt), log) # 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: if args.pretrained in PRETRAINED: # pretrained networks (downloaded automatically) print(">> Load pre-trained model from '{}'".format( PRETRAINED[args.pretrained])) state = load_url(PRETRAINED[args.pretrained], model_dir=os.path.join('data/networks')) 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']) else: print(">> Using pre-trained model on imagenet '{}'".format( args.arch)) model = init_network(model=args.arch, pooling=args.pool, whitening=args.whitening, pretrained=True) 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 = CustomizeTuplesDataset(name='7Scene', mode='train', db_file=args.db_file, ims_root=args.data_root, imsize=args.image_size, snum=args.sneg_num, dnum=args.dneg_num, qsize=args.query_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 = CustomizeTuplesDataset(name='7Scene', mode='val', db_file=args.db_file, ims_root=args.data_root, imsize=args.image_size, snum=args.sneg_num, dnum=args.dneg_num, qsize=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 gt_root = os.path.join(args.data_root, 'sfm_relative_pose_pairs') data_splits = split_dataset(args.data_root, test_datasets, val_step=6, seed=0) test(model, args.data_root, data_splits, gt_root, epoch=0, pass_thres=8, knn=10, query_key='val', db_key='train', log=log) 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() # train for one epoch on train set loss = train(train_loader, model, criterion, optimizer, epoch, print_freq=40, log=log) # evaluate on validation set if args.val and (epoch + 1) % 5 == 0: loss = validate(val_loader, model, criterion, epoch, print_freq=100, log=log) # evaluate on test datasets test(model, args.data_root, data_splits, gt_root, epoch, pass_thres=8, knn=10, query_key='val', db_key='train', log=log) # 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) print('Training Finished') log.close()
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)))
state = torch.load(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 net_params['multi_layer_cat'] = state['meta']['multi_layer_cat'] # 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 network_offtheshelf is not None: # parse off-the-shelf parameters offtheshelf = network_offtheshelf.split('-') net_params = {} net_params['architecture'] = offtheshelf[0]
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)) # 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() print (vecs.shape) print (qvecs.shape) # search, rank, and print scores = np.dot(vecs.T, qvecs) print (scores.shape) # to save scores (single query) # oxford #f = 'oxf_single.npy' # paris #f = 'par_single.npy' # roxford #f = 'roxf_single.npy' # rparis f = 'rpar_single.npy' 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) # save np.save(f, scores) ranks = np.argsort(-scores, axis=0) compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd']) """ ############################################################ # Test # 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'])] print(qimages) # to load scores # oxford #f = 'oxf_single.npy' #f = 'oxf_mq_avg.npy' #f = 'oxf_mq_max.npy' #f = 'oxf_sc_imf.npy' # paris #f = 'par_single.npy' #f = 'par_mq_avg.npy' #f = 'par_mq_max.npy' f = 'par_sc_imf.npy' # roxford #f = 'roxf_single.npy' #f = 'roxf_mq_avg.npy' #f = 'roxf_mq_max.npy' #f = 'roxf_sc_imf.npy' # rparis #f = 'rpar_single.npy' #f = 'rpar_mq_avg.npy' #f = 'rpar_mq_max.npy' #f = 'rpar_sc_imf.npy' # load scores = np.load(f) 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() # 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 start = time.time() print('>> {}: Extracting...') img_path = osp.join(args.data_path, "images", "images_upright") fnames_db = [ osp.join(path, fname) for path, _, files in os.walk(osp.join(img_path, "db")) for fname in files if fname[-3:].lower() in ["jpg", "png"] ] fnames_q = [ osp.join(path, fname) for path, _, files in os.walk(osp.join(img_path, "query")) for fname in files if fname[-3:].lower() in ["jpg", "png"] ] # extract vectors vecs_db = extract_vectors(net, fnames_db, args.image_size, transform, ms=ms) vecs_q = extract_vectors(net, fnames_q, args.image_size, transform, ms=ms) print('>> {}: Evaluating...') # convert to numpy vecs_db = vecs_db.numpy() vecs_q = vecs_q.numpy() # search, rank, and print scores_db = np.dot(vecs_db.T, vecs_db) scores_q = np.dot(vecs_db.T, vecs_q) ranks_db = np.argsort(-scores_db, axis=0) ranks_q = np.argsort(-scores_q, axis=0) print(ranks_q.shape, ranks_db.shape) images_db = [fname_db[len(img_path) + 1:] for fname_db in fnames_db] images_q = [fname_q[len(img_path) + 1:] for fname_q in fnames_q] pairs_db = [] for q_id in range(len(images_db)): img_q = images_db[q_id] pairs_per_q = [ " ".join([img_q, images_db[db_id]]) for db_id in list(ranks_db[1:args.top_n + 1, q_id]) ] pairs_db += pairs_per_q for q_id in range(len(images_q)): img_q = images_q[q_id] pairs_per_q = [ " ".join([img_q, images_db[db_id]]) for db_id in list(ranks_q[1:args.top_n + 1, q_id]) ] pairs_db += pairs_per_q with open( osp.join(args.data_path, "image_pairs_day_night_" + str(args.top_n) + ".txt"), "w") as f: for pair in pairs_db: f.write(pair + "\n") print("Done")
def main(): args = parser.parse_args() # 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: ") if "epoch" in state: print("Model after {} epochs".format(state["epoch"])) 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: test both single scale and multiscale ms_singlescale = [1] msp_singlescale = 1 ms_multiscale = list(eval(args.multiscale)) msp_multiscale = 1 if len(ms_multiscale ) > 1 and net.meta['pooling'] == 'gem' and not net.meta[ 'regional'] and not net.meta['whitening']: msp_multiscale = net.pool.p.item() print(">> Set-up multiscale:") print(">>>> ms: {}".format(ms_multiscale)) print(">>>> msp: {}".format(msp_multiscale)) # 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)) Lw = net.meta['Lw'][args.whitening] 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) 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: Lw = {} for whiten_type, ms, msp in zip( ["ss", "ms"], [ms_singlescale, ms_multiscale], [msp_singlescale, msp_multiscale]): print('>> {0}: Learning whitening {1}...'.format( args.whitening, whiten_type)) # 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[whiten_type] = {'m': m, 'P': P} print('>> {}: elapsed time: {}'.format( args.whitening, htime(time.time() - start))) # 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) else: Lw = None # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() for whiten_type, ms, msp in zip(["ss", "ms"], [ms_singlescale, ms_multiscale], [msp_singlescale, msp_multiscale]): print('>> Extracting feature on {0}, whitening {1}'.format( dataset, whiten_type)) # 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[whiten_type]['m'], Lw[whiten_type]['P']) qvecs_lw = whitenapply(qvecs, Lw[whiten_type]['m'], Lw[whiten_type]['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 {}'.format(whiten_type), 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]) # evaluate on test datasets datasets = args.datasets.split(',') for dataset in datasets: start = time.time() print('>> {}: Extracting...'.format(dataset)) print('>> Prepare data information...') index_file_path = os.path.join(get_data_root(), 'index.csv') index_mark_path = os.path.join(get_data_root(), 'index_mark.csv') index_miss_path = os.path.join(get_data_root(), 'index_miss.csv') test_file_path = os.path.join(get_data_root(), 'test.csv') test_mark_path = os.path.join(get_data_root(), 'test_mark.csv') test_mark_add_path = os.path.join(get_data_root(), 'test_mark_add.csv') test_miss_path = os.path.join(get_data_root(), 'test_miss.csv') if dataset == 'google-landmarks-dataset': index_img_path = os.path.join(get_data_root(), 'index') test_img_path = os.path.join(get_data_root(), 'google-landmarks-dataset-test') elif dataset == 'google-landmarks-dataset-resize': index_img_path = os.path.join(get_data_root(), 'resize_index_image') test_img_path = os.path.join(get_data_root(), 'resize_test_image') if not (os.path.isfile(index_mark_path) or os.path.isfile(index_miss_path)): clear_no_exist(index_file_path, index_mark_path, index_miss_path, index_img_path) if not (os.path.isfile(test_mark_path) or os.path.isfile(test_miss_path)): clear_no_exist(test_file_path, test_mark_path, test_miss_path, test_img_path) print('>> load index image path...') retrieval_other_dataset = '/home/iap205/Datasets/google-landmarks-dataset-resize' csvfile = open(index_mark_path, 'r') csvreader = csv.reader(csvfile) images = [] miss, add = 0, 0 for line in csvreader: if line[0] == '1': images.append(os.path.join(index_img_path, line[1] + '.jpg')) elif line[0] == '0': retrieval_img_path = os.path.join(retrieval_other_dataset, 'resize_index_image', line[1] + '.jpg') if os.path.isfile(retrieval_img_path): images.append(retrieval_img_path) add += 1 miss += 1 csvfile.close() print( '>>>> index image miss: {}, supplement: {}, still miss: {}'.format( miss, add, miss - add)) print('>> load query image path...') csvfile = open(test_mark_path, 'r') csvreader = csv.reader(csvfile) savefile = open(test_mark_add_path, 'w') save_writer = csv.writer(savefile) qimages = [] miss, add = 0, 0 for line in csvreader: if line[0] == '1': qimages.append(os.path.join(test_img_path, line[1] + '.jpg')) save_writer.writerow(line) elif line[0] == '0': retrieval_img_path = os.path.join(retrieval_other_dataset, 'resize_test_image', line[1] + '.jpg') if os.path.isfile(retrieval_img_path): qimages.append(retrieval_img_path) save_writer.writerow(['1', line[1]]) add += 1 else: save_writer.writerow(line) miss += 1 csvfile.close() savefile.close() print( '>>>> test image miss: {}, supplement: {}, still miss: {}'.format( miss, add, miss - add)) # extract index vectors print('>> {}: index images...'.format(dataset)) split_num = 6 extract_num = int(len(images) / split_num) num_list = list(range(0, len(images), extract_num)) num_list.append(len(images)) # k = 0 # print('>>>> extract part {} of {}'.format(k, split_num-1)) # vecs = extract_vectors(net, images[num_list[k]:num_list[k+1]], args.image_size, transform, ms=ms, msp=msp) # vecs = vecs.numpy() # print('>>>> save index vecs to pkl...') # vecs_file_path = os.path.join(get_data_root(), 'index_vecs{}_of_{}.pkl'.format(k+1, split_num)) # vecs_file = open(vecs_file_path, 'wb') # pickle.dump(vecs[:, num_list[k]:num_list[k+1]], vecs_file) # vecs_file.close() # print('>>>> index_vecs{}_of_{}.pkl save done...'.format(k+1, split_num)) for i in range(split_num): # vecs_temp = np.loadtxt(open(os.path.join(get_data_root(), 'index_vecs{}_of_{}.csv'.format(i+1, split_num)), "rb"), # delimiter=",", skiprows=0) with open( os.path.join( get_data_root(), 'index_vecs{}_of_{}.pkl'.format(i + 1, split_num)), 'rb') as f: vecs_temp = pickle.load(f) if i == 0: vecs = vecs_temp else: vecs = np.hstack((vecs, vecs_temp[:, :])) del vecs_temp gc.collect() print('\r>>>> index_vecs{}_of_{}.pkl load done...'.format( i + 1, split_num), end='') print('') # extract query vectors print('>> {}: query images...'.format(dataset)) split_num = 1 extract_num = int(len(qimages) / split_num) num_list = list(range(0, len(qimages), extract_num)) num_list.append(len(qimages)) # k = 0 # print('>>>> extract part {} of {}'.format(k, split_num - 1)) # qvecs = extract_vectors(net, qimages[num_list[k]:num_list[k + 1]], args.image_size, transform, ms=ms, msp=msp) # qvecs = qvecs.numpy() # for i in range(split_num): # qvecs_file_path = os.path.join(get_data_root(), 'test_vecs{}_of_{}.pkl'.format(i+1, split_num)) # qvecs_file = open(qvecs_file_path, 'wb') # pickle.dump(qvecs[:, num_list[i]:num_list[i+1]], qvecs_file) # qvecs_file.close() # print('\r>>>> test_vecs{}_of_{}.pkl save done...'.format(i+1, split_num), end='') # print('') for i in range(split_num): # qvecs_temp = np.loadtxt(open(os.path.join(get_data_root(), 'test_vecs{}_of_{}.csv'.format(i+1, split_num)), "rb"), # delimiter=",", skiprows=0) with open( os.path.join( get_data_root(), 'test_vecs{}_of_{}.pkl'.format(i + 1, split_num)), 'rb') as f: qvecs_temp = pickle.load(f) if i == 0: qvecs = qvecs_temp else: qvecs = np.hstack((qvecs, qvecs_temp[:, :])) del qvecs_temp gc.collect() print('\r>>>> test_vecs{}_of_{}.pkl load done...'.format( i + 1, split_num), end='') print('') # vecs = np.zeros((2048, 1093278)) # qvecs = np.zeros((2048, 115921)) # save vecs to csv file # np.savetxt(os.path.join(get_data_root(), 'index_vecs{}_of_{}.csv'.format(k, split_num-1)), vecs, delimiter=',') # np.savetxt(os.path.join(get_data_root(), 'test_vecs{}_of_{}.csv'.format(k, split_num-1)), qvecs, delimiter=',') # 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)) # extract whitening vectors print('>> {}: Extracting...'.format(args.whitening)) wvecs = vecs # learning whitening print('>> {}: Learning...'.format(args.whitening)) m, P = pcawhitenlearn(wvecs) 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 print('>> apply PCAwhiten...') if Lw is not None: # whiten the vectors and shrinkage vecs = whitenapply(vecs, Lw['m'], Lw['P']) qvecs = whitenapply(qvecs, Lw['m'], Lw['P']) print('>>>> save index PCAwhiten vecs to pkl...') split_num = 6 extract_num = int(len(images) / split_num) num_list = list(range(0, len(images), extract_num)) num_list.append(len(images)) for i in range(split_num): vecs_file_path = os.path.join( get_data_root(), 'index_PCAwhiten_vecs{}_of_{}.pkl'.format(i + 1, split_num)) vecs_file = open(vecs_file_path, 'wb') pickle.dump(vecs[:, num_list[i]:num_list[i + 1]], vecs_file) vecs_file.close() print( '\r>>>> index_PCAwhiten_vecs{}_of_{}.pkl save done...'.format( i + 1, split_num), end='') print('') print('>>>> save test PCAwhiten vecs to pkl...') split_num = 1 extract_num = int(len(qimages) / split_num) num_list = list(range(0, len(qimages), extract_num)) num_list.append(len(images)) for i in range(split_num): qvecs_file_path = os.path.join( get_data_root(), 'test_PCAwhiten_vecs{}_of_{}.pkl'.format(i + 1, split_num)) qvecs_file = open(qvecs_file_path, 'wb') pickle.dump(qvecs[:, num_list[i]:num_list[i + 1]], qvecs_file) qvecs_file.close() print('\r>>>> test_PCAwhiten_vecs{}_of_{}.pkl save done...'.format( i + 1, split_num), end='') print('') print('>>>> load index PCAwhiten vecs from pkl...') for i in range(split_num): with open( os.path.join( get_data_root(), 'index_PCAwhiten_vecs{}_of_{}.pkl'.format( i + 1, split_num)), 'rb') as f: vecs_temp = pickle.load(f) if i == 0: vecs = vecs_temp else: vecs = np.hstack((vecs, vecs_temp[:, :])) del vecs_temp gc.collect() print( '\r>>>> index_PCAwhiten_vecs{}_of_{}.pkl load done...'.format( i + 1, split_num), end='') print('') print('>>>> load test PCAwhiten vecs from pkl...') for i in range(split_num): with open( os.path.join( get_data_root(), 'test_PCAwhiten_vecs{}_of_{}.pkl'.format( i + 1, split_num)), 'rb') as f: qvecs_temp = pickle.load(f) if i == 0: qvecs = qvecs_temp else: qvecs = np.hstack((qvecs, qvecs_temp[:, :])) del qvecs_temp gc.collect() print('\r>>>> test_PCAwhiten_vecs{}_of_{}.pkl load done...'.format( i + 1, split_num), end='') print('') # extract principal components and dimension shrinkage ratio = 0.8 vecs = vecs[:int(vecs.shape[0] * ratio), :] qvecs = vecs[:int(qvecs.shape[0] * ratio), :] print('>> {}: Evaluating...'.format(dataset)) split_num = 50 top_num = 100 vecs_T = np.zeros((vecs.shape[1], vecs.shape[0])).astype('float32') vecs_T[:] = vecs.T[:] QE_iter = 0 QE_weight = (np.arange(top_num, 0, -1) / top_num).reshape( 1, top_num, 1) print('>> find {} nearest neighbour...'.format(top_num)) import faiss # place it in the file top will cause network load so slowly # ranks_top_100 = np.loadtxt(open(os.path.join(get_data_root(), 'ranks_top_{}.csv'.format(top_num)), "rb"), # delimiter=",", skiprows=0).astype('int') for iter in range(0, QE_iter + 1): if iter != 0: # ranks_top_100 = np.ones((100, 115921)).astype('int') print('>> Query expansion iteration {}'.format(iter)) ranks_split = 50 for i in range(ranks_split): ranks_top_100_split = ranks_top_100[:, int(ranks_top_100. shape[1] / ranks_split * i ):int(ranks_top_100 .shape[1] / ranks_split * (i + 1))] top_100_vecs = vecs[:, ranks_top_100_split] # (2048, 100, query_split_size) qvecs_temp = (top_100_vecs * QE_weight).sum(axis=1) qvecs_temp = qvecs_temp / (np.linalg.norm( qvecs_temp, ord=2, axis=0, keepdims=True) + 1e-6) if i == 0: qvecs = qvecs_temp else: qvecs = np.hstack((qvecs, qvecs_temp)) del ranks_top_100_split, top_100_vecs, qvecs_temp gc.collect() print('\r>>>> calculate new query vectors {}/{} done...'. format(i + 1, ranks_split), end='') print('') qe_iter_qvecs_path = os.path.join( get_data_root(), 'QE_iter{}_qvecs.pkl'.format(iter)) qe_iter_qvecs_file = open(qe_iter_qvecs_path, 'wb') pickle.dump(qvecs, qe_iter_qvecs_file) qe_iter_qvecs_file.close() print('>>>> QE_iter{}_qvecs.pkl save done...'.format(iter)) del ranks_top_100 gc.collect() for i in range(split_num): # scores = np.dot(vecs.T, qvecs[:, int(qvecs.shape[1]/split_num*i):int(qvecs.shape[1]/split_num*(i+1))]) # ranks = np.argsort(-scores, axis=0) # kNN search k = top_num index = faiss.IndexFlatL2(vecs.shape[0]) index.add(vecs_T) query_vecs = qvecs[:, int(qvecs.shape[1] / split_num * i):int(qvecs.shape[1] / split_num * (i + 1))] qvecs_T = np.zeros((query_vecs.shape[1], query_vecs.shape[0])).astype('float32') qvecs_T[:] = query_vecs.T[:] _, ranks = index.search(qvecs_T, k) ranks = ranks.T if i == 0: ranks_top_100 = ranks[:top_num, :] else: ranks_top_100 = np.hstack( (ranks_top_100, ranks[:top_num, :])) # del scores, ranks del index, query_vecs, qvecs_T, ranks gc.collect() print('\r>>>> kNN search {} nearest neighbour {}/{} done...'. format(top_num, i + 1, split_num), end='') del qvecs gc.collect() print('') del vecs, vecs_T gc.collect() # save to csv file print(">> save to submission.csv file...") submission_file = open(os.path.join(get_data_root(), 'submission.csv'), 'w') writer = csv.writer(submission_file) test_mark_file = open(test_mark_add_path, 'r') csvreader = csv.reader(test_mark_file) cnt = 0 writer.writerow(['id', 'images']) for index, line in enumerate(csvreader): (flag, img_name) = line[:2] if flag == '1': select = [] for i in range(top_num): select.append(images[int( ranks_top_100[i, cnt])].split('/')[-1].split('.jpg')[0]) cnt += 1 writer.writerow([ img_name.split('/')[-1].split('.jpg')[0], ' '.join(select) ]) else: # random_list = random.sample(range(0, len(images)), top_num) random_list = np.random.choice(len(images), top_num, replace=False) select = [] for i in range(top_num): select.append( images[random_list[i]].split('/')[-1].split('.jpg')[0]) writer.writerow([ img_name.split('/')[-1].split('.jpg')[0], ' '.join(select) ]) if cnt % 10 == 0 or cnt == len(qimages): print('\r>>>> {}/{} done...'.format(cnt, len(qimages)), end='') submission_file.close() test_mark_file.close() print('') print('>> {}: elapsed time: {}'.format(dataset, htime(time.time() - start)))
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() # 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()) # download_train('/media/hxq/ExtractedDatasets/Datasets/retrieval-SfM') # download_test('/media/hxq/ExtractedDatasets/Datasets/retrieval-SfM') # 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{}_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 # 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_params['multi_layer_cat'] = args.multi_layer_cat model = init_network(model_params) print(">>>> loaded model: ") print(model.meta_repr()) # 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, pool, whitening # IMPORTANT: no weight decay for pooling parameters 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 parameters p weight decay should be 0 parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0}) else: # regional, pooling parameters p weight decay should be 0, # and we want to add regional whitening if it is there parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, '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'] # hxq added, for whitening test # print('>> load the parameters of supervised whitening for the FC layer initialization') # whiten_fn = '/media/iap205/Data/Export/cnnimageretrieval-google_landmark_retrieval/trained_network/' \ # 'R101_O_GL_FC/google-landmarks-dataset-resize_resnet101_gem_whiten_contrastive_m0.85_' \ # 'adam_lr5.0e-07_wd1.0e-04_nnum5_qsize2000_psize22000_bsize5_imsize362/' \ # 'model_epoch114.pth.tar_google-landmarks-dataset_whiten_ms.pth' # Lw = torch.load(whiten_fn) # P = Lw['P'] # m = Lw['m'] # P = torch.from_numpy(P).float() # m = torch.from_numpy(m).float() # checkpoint['state_dict']['whiten.weight'] = P # checkpoint['state_dict']['whiten.bias'] = -torch.mm(P, m).squeeze() 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) # print(">> Use new scheduler") else: print(">> No checkpoint found at '{}'".format(args.resume)) # Data loading # hxq added, but it can be generate in advance # if args.training_dataset == 'google-landmarks-dataset-v2': # data_table_path = os.path.join(get_data_root(), 'train.csv') # img_check_path = os.path.join(get_data_root(), 'train') # train_val_file_output_path = os.path.join(get_data_root()) # test_file_output_path = os.path.join(get_data_root(), 'test', args.training_dataset+'-test') # val_set_size = 10000 # test_set_size = 10000 # elif args.training_dataset == 'google-landmarks-dataset-resize': # data_table_path = os.path.join(get_data_root(), 'train.csv') # img_check_path = os.path.join(get_data_root(), 'resize_train_image') # train_val_file_output_path = os.path.join(get_data_root()) # test_file_output_path = os.path.join(get_data_root(), 'test', args.training_dataset+'-test') # val_set_size = 10000 # test_set_size = 10000 # elif args.training_dataset == 'google-landmarks-dataset': # pass # if args.training_dataset == 'google-landmarks-dataset-v2'\ # or args.training_dataset == 'google-landmarks-dataset-resize'\ # or args.training_dataset == 'google-landmarks-dataset': # if not (os.path.isfile(os.path.join(train_val_file_output_path, '{}.pkl'.format(args.training_dataset))) # and os.path.isfile(os.path.join(test_file_output_path, 'gnd_{}-test.pkl'.format(args.training_dataset)))): # gen_train_val_test_pkl(args.training_dataset, data_table_path, img_check_path, val_set_size, test_set_size, # train_val_file_output_path, test_file_output_path) 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) # add by hxq # save_checkpoint({ # 'epoch': 0, # 'meta': model.meta, # 'state_dict': model.state_dict(), # 'min_loss': min_loss, # 'optimizer': optimizer.state_dict(), # }, 0, args.directory) # hxq added, save trained network as epoch # network_path = {'vgg16': '/home/iap205/Datasets/retrieval-SfM/networks/retrievalSfM120k-vgg16-gem-b4dcdc6.pth', # 'resnet101': '/home/iap205/Datasets/retrieval-SfM/networks/retrievalSfM120k-resnet101-gem-b80fb85.pth'} # model_trained = torch.load(network_path['resnet101']) # save_checkpoint({ # 'epoch': 0, # 'meta': model.meta, # 'state_dict': model_trained['state_dict'], # 'min_loss': min_loss, # 'optimizer': optimizer.state_dict(), # }, 0, args.directory) 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)