def eval_global(net, inference, globals, *, datasets): """Evaluate global descriptors""" net.eval() time0 = time.time() logger = globals["logger"] logger.info("Starting global evaluation") results = {} for dataset in datasets: images, qimages, bbxs, gnd = data_helpers.load_dataset(dataset, data_root=globals['root_path']) logger.info(f"Evaluating {dataset}") with logging.LoggingStopwatch("extracting database images", logger.info, logger.debug): dset = ImagesFromList(root='', images=images, imsize=inference['image_size'], bbxs=None, transform=globals['transform']) vecs = how_net.extract_vectors(net, dset, globals["device"], scales=inference['scales']) with logging.LoggingStopwatch("extracting query images", logger.info, logger.debug): qdset = ImagesFromList(root='', images=qimages, imsize=inference['image_size'], bbxs=bbxs, transform=globals['transform']) qvecs = how_net.extract_vectors(net, qdset, globals["device"], scales=inference['scales']) vecs, qvecs = vecs.numpy(), qvecs.numpy() ranks = np.argsort(-np.dot(vecs, qvecs.T), axis=0) results[dataset] = score_helpers.compute_map_and_log(dataset, ranks, gnd, logger=logger) logger.info(f"Finished global evaluation in {int(time.time()-time0) // 60} min") return results
def extract_vectors(net, images, image_size, transform, bbxs=None, ms=[1], msp=1, print_freq=10): # moving network to gpu and eval mode net.cuda() net.eval() # creating dataset loader loader = torch.utils.data.DataLoader(ImagesFromList(root='', images=images, imsize=image_size, bbxs=bbxs, transform=transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) # extracting vectors with torch.no_grad(): vecs = torch.zeros(net.meta['outputdim'], len(images)) for i, input in enumerate(loader): input = input.cuda() #input.shape = (1, 3, H, W) if len(ms) == 1: vecs[:, i] = extract_ss(net, input) else: # hxq modified # if max(input.shape[-1], input.shape[-2]) >= image_size: # vecs[:, i] = extract_ms(net, input, ms, msp) # else: # vecs[:, i] = extract_ss(net, input) vecs[:, i] = extract_ms(net, input, ms, msp) if (i + 1) % print_freq == 0 or (i + 1) == len(images): print('\r>>>> {}/{} done...'.format((i + 1), len(images)), end='') print('') # hxq modified, test add features map for retrieval # vecs = [] # for i, input in enumerate(loader): # input = input.cuda() # # if len(ms) == 1: # vecs.append(extract_ss(net, input)) # else: # vecs.append(extract_ms(net, input, ms, msp)) # # if (i+1) % print_freq == 0 or (i+1) == len(images): # print('\r>>>> {}/{} done...'.format((i+1), len(images)), end='') # print('') return vecs
def extract_vectors(net, images, image_size, transform, bbxs=None, ms=[1], msp=1, print_freq=10): # moving network to gpu and eval mode net.cuda() net.eval() # creating dataset loader loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=images, imsize=image_size, bbxs=bbxs, transform=transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) # extracting vectors with torch.no_grad(): vecs = torch.zeros(net.meta['outputdim'], len(images)) for i, input in enumerate(loader): input = input.cuda() if len(ms) == 1 and ms[0] == 1: vecs[:, i] = extract_ss(net, input) else: vecs[:, i] = extract_ms(net, input, ms, msp) if (i+1) % print_freq == 0 or (i+1) == len(images): print('\r>>>> {}/{} done...'.format((i+1), len(images)), end='') print('') return vecs
def extract_local_vectors(net, images, image_size, transform, bbxs=None, ms=[1], msp=1, print_freq=10): # moving network to gpu and eval mode net.cuda() net.eval() # creating dataset loader loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=images, imsize=image_size, bbxs=bbxs, transform=transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) # extracting vectors with torch.no_grad(): vecs = [] for i, input in enumerate(loader): input = input.cuda() if len(ms) == 1: vecs.append(extract_ssl(net, input)) else: # TODO: not implemented yet # vecs.append(extract_msl(net, input, ms, msp)) raise NotImplementedError if (i+1) % print_freq == 0 or (i+1) == len(images): print('\r>>>> {}/{} done...'.format((i+1), len(images)), end='') print('') return vecs
def extract_vectors(model, images, image_size=1024, eval_transform=None, bbxs=None, scales=(1, ), tta_gem_p=1.0, tqdm_desc=''): if eval_transform is None: eval_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) local_eval_loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=images, imsize=image_size, bbxs=bbxs, transform=eval_transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) ids, feats = [], [] for i, x in tqdm(enumerate(local_eval_loader), total=len(local_eval_loader), desc=tqdm_desc, miniters=None, ncols=55): batch_size, _, h, w = x.shape feat_blend = [] with torch.no_grad(): x = x.to('cuda') for s in scales: size = int(h * s // model.DIVIDABLE_BY * model.DIVIDABLE_BY), \ int(w * s // model.DIVIDABLE_BY * model.DIVIDABLE_BY) # round off scaled_x = F.interpolate(x, size=size, mode='bilinear', align_corners=True) feat = model.extract_feat(scaled_x) if tta_gem_p != 1.0: feat = feat**tta_gem_p feat = feat.cpu().numpy() feat_blend.append(feat) feat_blend = np.mean(feat_blend, axis=0) feats.append(feat_blend) feats = np.concatenate(feats) if tta_gem_p != 1.0: feats = feats**(1.0 / tta_gem_p) feats = utils.l2norm_numpy(feats) return feats
def asmk_index_database(net, inference, globals, logger, *, asmk, images, distractors_path=None): """Asmk evaluation step 'aggregate_database' and 'build_ivf'""" data_opts = {"imsize": inference['image_size'], "transform": globals['transform']} infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']} # Index database vectors dset = ImagesFromList(root='', images=images, bbxs=None, **data_opts) vecs, imids, *_ = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts) asmk_dataset = asmk.build_ivf(vecs, imids, distractors_path=distractors_path) logger.info(f"Indexed images in {asmk_dataset.metadata['build_ivf']['index_time']:.2f}s") logger.debug(f"IVF stats: {asmk_dataset.metadata['build_ivf']['ivf_stats']}") return asmk_dataset
def asmk_train_codebook(net, inference, globals, logger, *, codebook_training, asmk, cache_path): """Asmk evaluation step 'train_codebook'""" if cache_path and cache_path.exists(): return asmk.train_codebook(None, cache_path=cache_path) images = data_helpers.load_dataset('train', data_root=globals['root_path'])[0] images = images[:codebook_training['images']] dset = ImagesFromList(root='', images=images, imsize=inference['image_size'], bbxs=None, transform=globals['transform']) infer_opts = {"scales": codebook_training['scales'], "features_num": inference['features_num']} des_train = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts)[0] asmk = asmk.train_codebook(des_train, cache_path=cache_path) logger.info(f"Codebook trained in {asmk.metadata['train_codebook']['train_time']:.1f}s") return asmk
def initialize_dim_reduction(net, globals, *, images, features_num): """Initialize dimensionality reduction by PCA whitening from 'images' number of descriptors""" if not net.dim_reduction: return if features_num is None: features_num = net.runtime['features_num'] images = data_helpers.load_dataset('train', data_root=globals['root_path'])[0][:images] dataset = ImagesFromList(root='', images=images, imsize=net.runtime['image_size'], bbxs=None, transform=globals["transform"]) des_train = how_net.extract_vectors_local(net.copy_excluding_dim_reduction(), dataset, globals["device"], scales=net.runtime['training_scales'], features_num=features_num)[0] net.dim_reduction.initialize_pca_whitening(des_train)
def asmk_aggregate_database(net, inference, globals, logger, *, asmk, images, partition, cache_path): """Asmk evaluation step 'aggregate_database'""" if cache_path.exists(): logger.debug("Step results already exist") return codebook = asmk.codebook kernel = kern_pkg.ASMKKernel(codebook, **asmk.params['build_ivf']['kernel']) start, end = int(len(images)*partition['norm_start']), int(len(images)*partition['norm_end']) data_opts = {"imsize": inference['image_size'], "transform": globals['transform']} infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']} # Aggregate database dset = ImagesFromList(root='', images=images[start:end], bbxs=None, **data_opts) vecs, imids, *_ = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts) imids += start quantized = codebook.quantize(vecs, imids, **asmk.params["build_ivf"]["quantize"]) aggregated = kernel.aggregate(*quantized, **asmk.params["build_ivf"]["aggregate"]) with cache_path.open("wb") as handle: pickle.dump(dict(zip(["des", "word_ids", "image_ids"], aggregated)), handle)
def asmk_query_ivf(net, inference, globals, logger, *, dataset, asmk_dataset, qimages, bbxs, gnd, results, cache_path, imid_offset=0): """Asmk evaluation step 'query_ivf'""" if gnd is None and cache_path and cache_path.exists(): logger.debug("Step results already exist") return data_opts = {"imsize": inference['image_size'], "transform": globals['transform']} infer_opts = {"scales": inference['scales'], "features_num": inference['features_num']} # Query vectors qdset = ImagesFromList(root='', images=qimages, bbxs=bbxs, **data_opts) qvecs, qimids, *_ = how_net.extract_vectors_local(net, qdset, globals["device"], **infer_opts) qimids += imid_offset metadata, query_ids, ranks, scores = asmk_dataset.query_ivf(qvecs, qimids) logger.debug(f"Average query time (quant+aggr+search) is {metadata['query_avg_time']:.3f}s") # Evaluate if gnd is not None: results[dataset] = score_helpers.compute_map_and_log(dataset, ranks.T, gnd, logger=logger) with cache_path.open("wb") as handle: pickle.dump({"metadata": metadata, "query_ids": query_ids, "ranks": ranks, "scores": scores}, handle)
def extract_train_descriptors(net, globals, *, images, features_num): """Extract descriptors for a given number of images from the train set""" if features_num is None: features_num = net.runtime['features_num'] images = data_helpers.load_dataset( 'train', data_root=globals['root_path'])[0][:images] dataset = ImagesFromList(root='', images=images, imsize=net.runtime['image_size'], bbxs=None, transform=globals["transform"]) des_train = how_net.extract_vectors_local( net, dataset, globals["device"], scales=net.runtime['training_scales'], features_num=features_num)[0] return des_train
def extract_vectors_a(net, images, image_size, image_resize, transform, bbxs=None, print_freq=10): # moving network to gpu and eval mode net.cuda() net.eval() # creating dataset loader loader = torch.utils.data.DataLoader(ImagesFromList(root='', images=images, imsize=image_size, bbxs=bbxs, transform=transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) # extracting vectors with torch.no_grad(): vecs = torch.zeros(net.meta['outputdim'], len(images)) for i, input in enumerate(loader): input = input.cuda() input_t = nn.functional.interpolate(input, scale_factor=image_resize / image_size, mode='bilinear', align_corners=False) vecs[:, i] = net(input_t).cpu().data.squeeze() if (i + 1) % print_freq == 0 or (i + 1) == len(images): print('\r>>>> {}/{} done...'.format((i + 1), len(images)), end='') print('') return vecs
def extract_vectors(net, images, image_size, transform, bbxs=None, ms=[1], msp=1, print_freq=10): # moving network to gpu and eval mode if torch.cuda.is_available(): net.cuda() net.eval() # creating dataset loader loader = torch.utils.data.DataLoader(ImagesFromList(root='', images=images, imsize=image_size, bbxs=bbxs, transform=transform), batch_size=1, shuffle=False, num_workers=1, pin_memory=True) # extracting vectors with torch.no_grad(): vecs = torch.zeros(net.meta['outputdim'], len(images)) img_paths = list() for i, (input, path) in enumerate(loader): if torch.cuda.is_available(): input = input.cuda() if len(ms) == 1 and ms[0] == 1: vecs[:, i] = extract_ss(net, input) else: vecs[:, i] = extract_ms(net, input, ms, msp) img_paths.append(path) if (i + 1) % print_freq == 0 or (i + 1) == len(images): print('\r>>>> {}/{} done...'.format((i + 1), len(images)), end='') print('') # vecs = torch.zeros(net.meta['outputdim'], len(images)) # img_path_list = list() # for i in range(len(images)): # img_path = images[i] # img_path_list.append(img_path) # input = img2tensor(img_path, image_size, transform) # if torch.cuda.is_available(): # input = input.cuda() # # if len(ms) == 1 and ms[0] == 1: # vecs[:, i] = extract_ss(net, input) # else: # vecs[:, i] = extract_ms(net, input, ms, msp) # # if (i + 1) % print_freq == 0 or (i + 1) == len(images): # print('\r>>>> {}/{} done...'.format((i + 1), len(images)), end='') return vecs, img_paths
def create_epoch_tuples(self, net, temp_loss): print('>> Creating tuples for an epoch of {}-{}...'.format( self.name, self.mode)) ## ------------------------ ## SELECTING POSITIVE PAIRS ## ------------------------ # draw qsize random queries for tuples idxs2qpool = torch.randperm(len(self.qpool))[:self.qsize] self.qidxs = [self.qpool[i] for i in idxs2qpool] self.pidxs = [self.ppool[i] for i in idxs2qpool] ## ------------------------ ## SELECTING NEGATIVE PAIRS ## ------------------------ # if nnum = 0 create dummy nidxs # useful when only positives used for training if self.nnum == 0: self.nidxs = [[] for _ in range(len(self.qidxs))] return 0 # draw poolsize random images for pool of negatives images ####if using less data newimages = [] for i in range(len(self.images)): if self.clusters[i] != -1: newimages.append(i) random.shuffle(newimages) idxs2images = newimages[:self.poolsize] # prepare network net.cuda() net.eval() batchsize = min(100, self.qsize) # no gradients computed, to reduce memory and increase speed with torch.no_grad(): print('>> Extracting descriptors for query images...') # prepare query loader loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in self.qidxs], imsize=self.imsize, re_size=True, transform=self.transform), batch_size=batchsize, shuffle=False, num_workers=8, pin_memory=True) # extract query vectors qvecs = torch.zeros(net.meta['outputdim'], len(self.qidxs)).cuda() for i, input in enumerate(loader): if batchsize != 1: if (i + 1) != len(loader): qvecs[:, i * batchsize:(i + 1) * batchsize] = net( input.cuda()).data.squeeze() else: qvecs[:, i * batchsize:len(self.qidxs)] = net( input.cuda()).data.squeeze() else: qvecs[:, i] = net(input.cuda()).data.squeeze() if (i + 1) % self.print_freq == 0 or (i + 1) == len(loader): print('\r>>>> {}/{} done...'.format((i + 1), len(loader)), end='') print('') print('>> Extracting descriptors for negative pool...') # prepare negative pool data loader loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in idxs2images], imsize=self.imsize, re_size=True, transform=self.transform), batch_size=batchsize, shuffle=False, num_workers=8, pin_memory=True) # extract negative pool vectors poolvecs = torch.zeros(net.meta['outputdim'], len(idxs2images)).cuda() for i, input in enumerate(loader): if batchsize != 1: if (i + 1) != len(loader): poolvecs[:, i * batchsize:(i + 1) * batchsize] = net( input.cuda()).data.squeeze() else: poolvecs[:, i * batchsize:len(idxs2images)] = net( input.cuda()).data.squeeze() else: poolvecs[:, i] = net(input.cuda()).data.squeeze() if (i + 1) % self.print_freq == 0 or (i + 1) == len(loader): print('\r>>>> {}/{} done...'.format((i + 1), len(loader)), end='') # poolvecs[:, i] = net(input.cuda()).data.squeeze() # if (i + 1) % self.print_freq == 0 or (i + 1) == len(idxs2images): # print('\r>>>> {}/{} done...'.format(i + 1, len(idxs2images)), end='') print('') print('>> Searching for hard negatives...') torch.cuda.empty_cache() # compute dot product scores and ranks on GPU scores = torch.mm(poolvecs.t(), qvecs) if self.using_cdvs != 0: print('using global descriptor') qvecs_global = torch.from_numpy( np.vstack([self.global_cdvs[:, i] for i in self.qidxs ])).float().cuda().transpose(1, 0) # qvecs = torch.cat((qvecs_global, qvecs), 0) poolvecs_global = torch.from_numpy( np.vstack([self.global_cdvs[:, i] for i in idxs2images ])).float().cuda().transpose(1, 0) # poolvecs = torch.cat((poolvecs_global, poolvecs), 0) scores_global = torch.mm(poolvecs_global.t(), qvecs_global) if self.using_cdvs != -1: scores = scores + scores_global * self.using_cdvs else: # factor=(torch.exp(torch.tensor(-0.55/temp_loss))) factor = temp_loss print("global param:{}".format(factor)) scores = scores + scores_global * factor scores, ranks = torch.sort(scores, dim=0, descending=True) avg_ndist = torch.tensor(0).float().cuda() # for statistics n_ndist = torch.tensor(0).float().cuda() # for statistics # selection of negative examples self.nidxs = [] for q in range(len(self.qidxs)): # do not use query cluster, # those images are potentially positive qcluster = self.clusters[self.qidxs[q]] clusters = [qcluster] nidxs = [] r = 0 while len(nidxs) < self.nnum: potential = idxs2images[ranks[r, q]] # take at most one image from the same cluster if (not self.clusters[potential] in clusters) and (self.clusters[potential] != -1): nidxs.append(potential) clusters.append(self.clusters[potential]) avg_ndist += torch.pow( qvecs[:, q] - poolvecs[:, ranks[r, q]] + 1e-6, 2).sum(dim=0).sqrt() n_ndist += 1 r += 1 self.nidxs.append(nidxs) print('>>>> Average negative l2-distance: {:.2f}'.format( avg_ndist / n_ndist)) print('>>>> Done') return (avg_ndist / n_ndist).item() # return average negative l2-distance
def extract_positive_vectors(self, net, target_data_idxs=[], save_embeds=False, save_embeds_epoch=-1, save_embeds_step=-1, save_embeds_total_steps=-1, save_embeds_path=''): # prepare network net.cuda() # if net was in training mode, temporarily switch to eval mode was_training = net.training if was_training: net.eval() # no gradients computed, to reduce memory and increase speed with torch.no_grad(): if len(target_data_idxs) > 0: # Refresh just a single data point specified by target_data_index pidxs = [self.pidxs[t] for t in target_data_idxs] else: # Rebuild all positive images within the dataset target_data_idxs = list(range(len(self.pidxs))) pidxs = self.pidxs images_to_rebuild = [self.images[i] for i in pidxs] print('>> Extracting descriptors for positive images...') # prepare positive image loader loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=images_to_rebuild, imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) assert len(loader) == len(target_data_idxs) # extract positive image vectors if self.pvecs is None: self.pvecs = torch.zeros(net.meta['outputdim'], len(self.pidxs)).cuda() j = 1 for i, image in zip(target_data_idxs, loader): self.pvecs[:, i] = net(image.cuda()).data.squeeze() print('\r>>>> {}/{} done...'.format(j, len(target_data_idxs)), end='') j = j + 1 print('') # Serialize the positive vectors if save_embeds: print( ">>>>> Epoch {} Step {}/{} positive embeddings serialization start.".format(save_embeds_epoch, save_embeds_step, save_embeds_total_steps)) torch.save( self.pvecs, os.path.join(save_embeds_path, '{}_positive.pt'.format(save_embeds_step))) print( ">>>>> Epoch {} Step {}/{} positive embeddings serialization complete!".format(save_embeds_epoch, save_embeds_step, save_embeds_total_steps)) print() # Restore the training mode if was_training: net.train() net.apply(set_batchnorm_eval)
def create_epoch_tuples(self, net): print('>> Creating tuples for an epoch of {}-{}...'.format('7Scene', self.mode)) # Draw qsize random queries for tuples idxs2qpool = torch.randperm(len(self.qpool))[:self.qsize] self.qidxs = [self.qpool[i] for i in idxs2qpool] self.pidxs = [self.ppool[i] for i in idxs2qpool] # Create dummy nidxs useful when only positives used for training if (self.snum + self.dnum) == 0: self.nidxs = [[] for _ in range(len(self.qidxs))] return # Extract descriptors for query images net.cuda() net.eval() t1 = time.time() loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=self.qims, imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) qimvecs = torch.Tensor(net.meta['outputdim'], len(self.qims)).cuda() for i, input in enumerate(loader): print('\r>>>>Extracting desc of query ims: {}/{} done...'.format(i+1, len(self.qims)), end='') qimvecs[:, i] = net(Variable(input.cuda())).data.squeeze() #print('Total time from extracting query image descriptors: {}s, scriptor size {} gpu memory usage {:.2f}%'.format(time.time()-t1, qimvecs.size(), get_gpu_mem_usage())) # Extract descriptors for db images if self.qkey == self.dbkey: # Query images are the same as db images dbimvecs = qimvecs else: t1 = time.time() loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=self.dbims, imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) dbimvecs = torch.Tensor(net.meta['outputdim'], len(self.dbims)).cuda() for i, input in enumerate(loader): print('\r>>>>Extracting desc of db ims {}/{} done...'.format(i+1, len(self.dbims)), end='') dbimvecs[:, i] = net(Variable(input.cuda())).data.squeeze() #print('Total time from extracting db image descriptors: {}s, scriptor size {} gpu memory usage {:.2f}%'.format(time.time()-t1, dbimvecs.size(), get_gpu_mem_usage())) # Search for negative pairs print('Searching negative pairs: snn {} dnn {}'.format(self.snum, self.dnum)) t1 = time.time() self.nidxs = [] avg_ndist = torch.Tensor([0]).cuda() n_ndist = torch.Tensor([0]).cuda() for qid in self.qidxs: qvecs = qimvecs[:, qid] qcid = self.qcids[qid] # Pick snum hard negative images from the same cluster nnpool = self.snn[qid] snvecs = dbimvecs[:, nnpool] scores = torch.mv(snvecs.t(), qvecs) scores, ranks = torch.sort(scores, dim=0, descending=True) snid = [] r = 0 while len(snid) < self.snum: # Ignore the self frame psnid = nnpool[ranks[r]] r += 1 if self.qkey == self.dbkey and psnid == qid: continue snid.append(psnid) avg_ndist += torch.pow(qvecs - dbimvecs[:, psnid] + 1e-6, 2).sum(dim=0).sqrt() n_ndist += 1 # Pick any dnum negatives from a different cluster dnid = [] dcid = [] while len(dnid) < self.dnum: pdnid = int(torch.randint(low=0, high=len(self.dbims), size=(1,))) pdcid = self.dbcids[pdnid] if pdcid == qcid or pdcid in dcid: # A different cluster from the query and selected negative ims continue dnid.append(pdnid) dcid.append(pdcid) avg_ndist += torch.pow(qvecs - dbimvecs[:, pdnid] + 1e-6, 2).sum(dim=0).sqrt() n_ndist += 1 self.nidxs.append(snid + dnid) avg_dist = list((avg_ndist/n_ndist).cpu())[0] #print('>>>> Average negative distance: {:.2f}'.format(avg_dist)) print('Total search time {}s'.format(time.time() - t1)) print('>>>> Dataset Preparation Done') return avg_dist
def create_epoch_tuples(self, net): print('>> Creating tuples for an epoch of {}-{}...'.format( self.name, self.mode)) ## ------------------------ ## SELECTING POSITIVE PAIRS ## ------------------------ # draw qsize random queries for tuples idxs2qpool = torch.randperm(len(self.qpool))[:self.qsize] self.qidxs = [self.qpool[i] for i in idxs2qpool] self.pidxs = [self.ppool[i] for i in idxs2qpool] ## ------------------------ ## SELECTING NEGATIVE PAIRS ## ------------------------ # if nnum = 0 create dummy nidxs # useful when only positives used for training if self.nnum == 0: self.nidxs = [[] for _ in range(len(self.qidxs))] return # draw poolsize random images for pool of negatives images idxs2images = torch.randperm(len(self.images))[:self.poolsize] # prepare network net.cuda() net.eval() print('>> Extracting descriptors for query images...') loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in self.qidxs], imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) qvecs = torch.Tensor(net.meta['outputdim'], len(self.qidxs)).cuda() for i, input in enumerate(loader): print('\r>>>> {}/{} done...'.format(i + 1, len(self.qidxs)), end='') qvecs[:, i] = net(Variable( input.cuda())).data.squeeze() ### squeeze is added print('') print('>> Extracting descriptors for negative pool...') loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in idxs2images], imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) poolvecs = torch.Tensor(net.meta['outputdim'], len(idxs2images)).cuda() for i, input in enumerate(loader): print('\r>>>> {}/{} done...'.format(i + 1, len(idxs2images)), end='') poolvecs[:, i] = net(Variable( input.cuda())).data.squeeze() ### squeeze added print('') print('>> Searching for hard negatives...') scores = torch.mm(poolvecs.t(), qvecs) scores, ranks = torch.sort(scores, dim=0, descending=True) self.nidxs = [] for q in range(len(self.qidxs)): qcluster = self.clusters[self.qidxs[q]] clusters = [qcluster] nidxs = [] r = 0 avg_ndist = torch.Tensor([0]).cuda() n_ndist = torch.Tensor([0]).cuda() while len(nidxs) < self.nnum: potential = idxs2images[ranks[r, q]] # take at most one image from the same cluster if not self.clusters[potential] in clusters: nidxs.append(potential) clusters.append(self.clusters[potential]) avg_ndist += torch.pow( qvecs[:, q] - poolvecs[:, ranks[r, q]] + 1e-6, 2).sum(dim=0).sqrt() n_ndist += 1 r += 1 self.nidxs.append(nidxs) print('>>>> Average negative distance: {:.2f}'.format( list((avg_ndist / n_ndist).cpu())[0])) print('>>>> Done')
def create_epoch_tuples(self, net): print('>> Creating tuples for an epoch of {}-{}...'.format( self.name, self.mode)) print(">>>> used network: ") print(net.meta_repr()) ## ------------------------ ## SELECTING POSITIVE PAIRS ## ------------------------ # draw qsize random queries for tuples idxs2qpool = torch.randperm(len(self.qpool))[:self.qsize] self.qidxs = [self.qpool[i] for i in idxs2qpool] self.pidxs = [self.ppool[i] for i in idxs2qpool] ## ------------------------ ## SELECTING NEGATIVE PAIRS ## ------------------------ # if nnum = 0 create dummy nidxs # useful when only positives used for training if self.nnum == 0: self.nidxs = [[] for _ in range(len(self.qidxs))] return 0 # draw poolsize random images for pool of negatives images idxs2images = torch.randperm(len(self.images))[:self.poolsize] # prepare network net.cuda() net.eval() # no gradients computed, to reduce memory and increase speed with torch.no_grad(): print('>> Extracting descriptors for query images...') # prepare query loader loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in self.qidxs], imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) # extract query vectors # meta['outputdim'] ? qvecs = torch.zeros(net.meta['outputdim'], len(self.qidxs)).cuda() for i, input in enumerate(loader): qvecs[:, i] = net(input.cuda()).data.squeeze() if (i + 1) % self.print_freq == 0 or (i + 1) == len( self.qidxs): print('\r>>>> {}/{} done...'.format( i + 1, len(self.qidxs)), end='') print('') print('>> Extracting descriptors for negative pool...') # prepare negative pool data loader loader = torch.utils.data.DataLoader(ImagesFromList( root='', images=[self.images[i] for i in idxs2images], imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True) # extract negative pool vectors poolvecs = torch.zeros(net.meta['outputdim'], len(idxs2images)).cuda() for i, input in enumerate(loader): poolvecs[:, i] = net(input.cuda()).data.squeeze() if (i + 1) % self.print_freq == 0 or (i + 1) == len(idxs2images): print('\r>>>> {}/{} done...'.format( i + 1, len(idxs2images)), end='') print('') print('>> Searching for hard negatives...') # compute dot product scores and ranks on GPU # poolvecs.t() ?? scores = torch.mm(poolvecs.t(), qvecs) scores, ranks = torch.sort(scores, dim=0, descending=True) avg_ndist = torch.tensor(0).float().cuda() # for statistics n_ndist = torch.tensor(0).float().cuda() # for statistics # selection of negative examples self.nidxs = [] for q in range(len(self.qidxs)): # do not use query cluster, # those images are potentially positive qcluster = self.clusters[self.qidxs[q]] clusters = [qcluster] nidxs = [] r = 0 while len(nidxs) < self.nnum: potential = idxs2images[ranks[r, q]] # take at most one image from the same cluster if not self.clusters[potential] in clusters: nidxs.append(potential) clusters.append(self.clusters[potential]) avg_ndist += torch.pow( qvecs[:, q] - poolvecs[:, ranks[r, q]] + 1e-6, 2).sum(dim=0).sqrt() n_ndist += 1 r += 1 self.nidxs.append(nidxs) print('>>>> Average negative l2-distance: {:.2f}'.format( avg_ndist / n_ndist)) print('>>>> Done') return (avg_ndist / n_ndist).item() # return average negative l2-distance
def cluster(self, net): self.pidxs = [] self.nidxs = [] # draw poolsize random images for pool of negatives images idxs2images = torch.randperm(len(self.images))[:self.poolsize] # prepare network self.net = net net.cuda() net.eval() Lw = net.meta["Lw"]["retrieval-SfM-120k"]["ss"] Lw_m = torch.from_numpy(Lw["m"]).cuda().float() Lw_p = torch.from_numpy(Lw["P"]).cuda().float() print(">> Extracting descriptors for pool...") loader = torch.utils.data.DataLoader( ImagesFromList( root="", images=[self.images[i] for i in idxs2images], imsize=self.imsize, transform=self.transform, random=True, ), batch_size=1, shuffle=False, num_workers=8, pin_memory=True, ) fname = self.filename if os.path.exists(fname): self.poolvecs = torch.load(fname).cuda() else: self.poolvecs = torch.Tensor(net.meta["outputdim"], len(idxs2images) * TIMES).cuda() with torch.no_grad(): for _ in range(TIMES): print(_) for i, input in enumerate(loader): print( "\r>>>> {}/{} done...".format( i + 1, len(idxs2images)), end="", ) input = (input - self.norm[0]) / self.norm[1] b = net(Variable(input.cuda())) b = do_whiten(b, Lw_m, Lw_p) self.poolvecs[:, i + _ * len(idxs2images)] = b.squeeze() torch.save(self.poolvecs.cpu(), fname) print("") print(">> KMeans...") poolvecs = self.poolvecs.cpu().numpy().T fname = fname + ".KMeans" kmeans = KMeans(n_clusters=512, n_jobs=-1) kmeans.fit(poolvecs) clustered_pool = [] self.pool_clusters_centers = torch.from_numpy(kmeans.cluster_centers_) for i in range(kmeans.cluster_centers_.shape[0]): clustered_pool.append(poolvecs[kmeans.labels_ == i, :]) with open(fname, "wb") as f: pickle.dump({ "kmeans": kmeans, "clustered_pool": clustered_pool }, f)
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 extract_negative_pool_vectors(self, net, target_data_idxs=[], refresh_nidxs_vectors=True, refresh_nidxs_others_vectors=False, save_embeds=False, save_embeds_epoch=-1, save_embeds_step=-1, save_embeds_total_steps=-1, save_embeds_path=''): # prepare network net.cuda() # if net was in training mode, temporarily switch to eval mode was_training = net.training if was_training: net.eval() # no gradients computed, to reduce memory and increase speed with torch.no_grad(): if len(target_data_idxs) > 0: # Refresh just a single data point specified by target_data_index idxs2images = set() for idx in target_data_idxs: if refresh_nidxs_others_vectors: idxs2images = idxs2images.union(set([im_index.item() for im_index in self.nidxs_others[idx]])) if refresh_nidxs_vectors: idxs2images = idxs2images.union(set([im_index.item() for im_index in self.nidxs[idx]])) print("Negative pool rebuild - idxs2images:", str(idxs2images)) # Since self.poolvecs is created with self.idx2images, # we need to determine the exact locations within # self.idx2images to where the indexes in the new idxs2images # are located target_data_idxs = [torch.nonzero(self.idxs2images == im_index, as_tuple=False).squeeze().item() for im_index in idxs2images] else: # Rebuild all queries within the negative image pool target_data_idxs = list(range(len(self.idxs2images))) idxs2images = self.idxs2images if len(idxs2images) > 0: images_to_rebuild = [self.images[i] for i in idxs2images] print('>> Extracting descriptors for negative pool...') # prepare negative pool data loader loader = torch.utils.data.DataLoader( ImagesFromList(root='', images=images_to_rebuild, imsize=self.imsize, transform=self.transform), batch_size=1, shuffle=False, num_workers=8, pin_memory=True ) assert len(loader) == len(target_data_idxs) # extract negative pool vectors if self.poolvecs is None: self.poolvecs = torch.zeros(net.meta['outputdim'], len(self.idxs2images)).cuda() j = 1 for i, image in zip(target_data_idxs, loader): self.poolvecs[:, i] = net(image.cuda()).data.squeeze() print('\r>>>> {}/{} done...'.format(j, len(target_data_idxs)), end='') j = j + 1 print('') # Serialize the query vectors if save_embeds: print( ">>>>> Epoch {} Step {}/{} pool embeddings serialization start.".format(save_embeds_epoch, save_embeds_step, save_embeds_total_steps)) torch.save( self.poolvecs, os.path.join(save_embeds_path, '{}_pools.pt'.format(save_embeds_step))) print( ">>>>> Epoch {} Step {}/{} pool embeddings serialization complete!".format(save_embeds_epoch, save_embeds_step, save_embeds_total_steps)) print() # Restore the training mode if was_training: net.train() net.apply(set_batchnorm_eval) # Return how many images were actually reembedded return len(idxs2images)
def eval_asmk(net, inference, globals, *, datasets, codebook_training, asmk): """Evaluate local descriptors with ASMK""" net.eval() time0 = time.time() logger = globals["logger"] logger.info("Starting asmk evaluation") # Train codebook images = data_helpers.load_dataset('train', data_root=globals['root_path'])[0] images = images[:codebook_training['images']] dset = ImagesFromList(root='', images=images, imsize=inference['image_size'], bbxs=None, transform=globals['transform']) infer_opts = { "scales": codebook_training['scales'], "features_num": inference['features_num'] } des_train = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts)[0] asmk = ASMKMethod.initialize_untrained(asmk).train_codebook(des_train) logger.info( f"Codebook trained in {asmk.metadata['train_codebook']['train_time']:.1f}s" ) scores = {} for dataset in datasets: images, qimages, bbxs, gnd = data_helpers.load_dataset( dataset, data_root=globals['root_path']) data_opts = { "imsize": inference['image_size'], "transform": globals['transform'] } infer_opts = { "scales": inference['scales'], "features_num": inference['features_num'] } logger.info(f"Evaluating {dataset}") # Database vectors dset = ImagesFromList(root='', images=images, bbxs=None, **data_opts) des = how_net.extract_vectors_local(net, dset, globals["device"], **infer_opts) asmk_dataset = asmk.build_ivf(des[0], des[1]) logger.info( f"Indexed images in {asmk_dataset.metadata['build_ivf']['index_time']:.2f}s" ) logger.debug( f"IVF stats: {asmk_dataset.metadata['build_ivf']['ivf_stats']}") # Query vectors qdset = ImagesFromList(root='', images=qimages, bbxs=bbxs, **data_opts) qdes = how_net.extract_vectors_local(net, qdset, globals["device"], **infer_opts) metadata, _images, ranks, _scores = asmk_dataset.query_ivf( qdes[0], qdes[1]) logger.debug( f"Average query time (quant+aggr+search) is {metadata['query_avg_time']:.3f}s" ) scores[dataset] = score_helpers.compute_map_and_log(dataset, ranks.T, gnd, logger=logger) logger.info( f"Finished asmk evaluation in {int(time.time()-time0) // 60} min") return scores