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
