dmlc--dgl
Tong He
9c41c22d7b
* add hilander model implementation draft * use focal loss * fix * change data root * add necessary scripts * update download links * update * update example table * fix * update readme with numbers * add empty folder * only eval at the end * set up hilander * inform results may fluctuate * address comments Co-authored-by: sneakerkg <xiaotj1990327@gmail.com> Co-authored-by: Ubuntu <ubuntu@ip-172-31-19-212.us-east-2.compute.internal>
107 行
3.4 KiB
Python
107 行
3.4 KiB
Python
"""
|
|
This file re-uses implementation from https://github.com/yl-1993/learn-to-cluster
|
|
"""
|
|
import gc
|
|
from tqdm import tqdm
|
|
|
|
from .faiss_gpu import faiss_search_approx_knn
|
|
|
|
__all__ = ['faiss_search_knn']
|
|
|
|
def precise_dist(feat, nbrs, num_process=4, sort=True, verbose=False):
|
|
import torch
|
|
feat_share = torch.from_numpy(feat).share_memory_()
|
|
nbrs_share = torch.from_numpy(nbrs).share_memory_()
|
|
dist_share = torch.zeros_like(nbrs_share).float().share_memory_()
|
|
|
|
precise_dist_share_mem(feat_share,
|
|
nbrs_share,
|
|
dist_share,
|
|
num_process=num_process,
|
|
sort=sort,
|
|
verbose=verbose)
|
|
|
|
del feat_share
|
|
gc.collect()
|
|
return dist_share.numpy(), nbrs_share.numpy()
|
|
|
|
def precise_dist_share_mem(feat,
|
|
nbrs,
|
|
dist,
|
|
num_process=16,
|
|
sort=True,
|
|
process_unit=4000,
|
|
verbose=False):
|
|
from torch import multiprocessing as mp
|
|
num, _ = feat.shape
|
|
num_per_proc = int(num / num_process) + 1
|
|
|
|
for pi in range(num_process):
|
|
sid = pi * num_per_proc
|
|
eid = min(sid + num_per_proc, num)
|
|
|
|
kwargs={'feat': feat,
|
|
'nbrs': nbrs,
|
|
'dist': dist,
|
|
'sid': sid,
|
|
'eid': eid,
|
|
'sort': sort,
|
|
'process_unit': process_unit,
|
|
'verbose': verbose,
|
|
}
|
|
bmm(**kwargs)
|
|
|
|
def bmm(feat,
|
|
nbrs,
|
|
dist,
|
|
sid,
|
|
eid,
|
|
sort=True,
|
|
process_unit=4000,
|
|
verbose=False):
|
|
import torch
|
|
_, cols = dist.shape
|
|
batch_sim = torch.zeros((eid - sid, cols), dtype=torch.float32)
|
|
for s in tqdm(range(sid, eid, process_unit),
|
|
desc='bmm',
|
|
disable=not verbose):
|
|
e = min(eid, s + process_unit)
|
|
query = feat[s:e].unsqueeze(1)
|
|
gallery = feat[nbrs[s:e]].permute(0, 2, 1)
|
|
batch_sim[s - sid:e - sid] = torch.clamp(torch.bmm(query, gallery).view(-1, cols), 0.0, 1.0)
|
|
|
|
if sort:
|
|
sort_unit = int(1e6)
|
|
batch_nbr = nbrs[sid:eid]
|
|
for s in range(0, batch_sim.shape[0], sort_unit):
|
|
e = min(s + sort_unit, eid)
|
|
batch_sim[s:e], indices = torch.sort(batch_sim[s:e],
|
|
descending=True)
|
|
batch_nbr[s:e] = torch.gather(batch_nbr[s:e], 1, indices)
|
|
nbrs[sid:eid] = batch_nbr
|
|
dist[sid:eid] = 1. - batch_sim
|
|
|
|
def faiss_search_knn(feat,
|
|
k,
|
|
nprobe=128,
|
|
num_process=4,
|
|
is_precise=True,
|
|
sort=True,
|
|
verbose=False):
|
|
|
|
dists, nbrs = faiss_search_approx_knn(query=feat,
|
|
target=feat,
|
|
k=k,
|
|
nprobe=nprobe,
|
|
verbose=verbose)
|
|
|
|
if is_precise:
|
|
print('compute precise dist among k={} nearest neighbors'.format(k))
|
|
dists, nbrs = precise_dist(feat,
|
|
nbrs,
|
|
num_process=num_process,
|
|
sort=sort,
|
|
verbose=verbose)
|
|
|
|
return dists, nbrs
|