in python/dglke/models/ke_model.py [0:0]
def _embed_sim(self, head, tail, emb, sfunc='cosine', bcast=False, pair_ws=False, topk=10):
batch_size=DEFAULT_INFER_BATCHSIZE
if head is None:
head = th.arange(0, emb.shape[0])
else:
head = th.tensor(head)
if tail is None:
tail = th.arange(0, emb.shape[0])
else:
tail = th.tensor(tail)
head_emb = emb[head]
tail_emb = emb[tail]
if sfunc == 'cosine':
sim_func = cosine_dist
elif sfunc == 'l2':
sim_func = l2_dist
elif sfunc == 'l1':
sim_func = l1_dist
elif sfunc == 'dot':
sim_func = dot_dist
elif sfunc == 'ext_jaccard':
sim_func = extended_jaccard_dist
if pair_ws is True:
result = []
# chunked cal score
score = []
num_head = head.shape[0]
num_tail = tail.shape[0]
# calculating scores using mini-batch, the default batchsize if 1024
# This can avoid OOM when using GPU
for i in range((num_head + batch_size - 1) // batch_size):
sh_emb = head_emb[i * batch_size : (i + 1) * batch_size \
if (i + 1) * batch_size < num_head \
else num_head]
sh_emb = sh_emb.to(self._device)
st_emb = tail_emb[i * batch_size : (i + 1) * batch_size \
if (i + 1) * batch_size < num_head \
else num_head]
st_emb = st_emb.to(self._device)
score.append(sim_func(sh_emb, st_emb, pw=True).to(th.device('cpu')))
score = th.cat(score, dim=0)
topk_score, topk_sidx = th.topk(score,
k=topk if score.shape[0] > topk else score.shape[0],
dim=0)
sidx = th.argsort(topk_score, dim=0, descending=True)
sidx = topk_sidx[sidx]
score = score[sidx]
result.append((head[sidx],
tail[sidx],
score))
else:
num_head = head.shape[0]
num_tail = tail.shape[0]
# calculating scores using mini-batch, the default batchsize if 1024
# This can avoid OOM when using GPU
score = []
for i in range((num_head + batch_size - 1) // batch_size):
sh_emb = head_emb[i * batch_size : (i + 1) * batch_size \
if (i + 1) * batch_size < num_head \
else num_head]
sh_emb = sh_emb.to(self._device)
s_score = []
for j in range((num_tail + batch_size - 1) // batch_size):
st_emb = tail_emb[j * batch_size : (j + 1) * batch_size \
if (j + 1) * batch_size < num_tail \
else num_tail]
st_emb = st_emb.to(self._device)
s_score.append(sim_func(sh_emb, st_emb).to(th.device('cpu')))
score.append(th.cat(s_score, dim=1))
score = th.cat(score, dim=0)
if bcast is False:
result = []
idx = th.arange(0, num_head * num_tail)
score = th.reshape(score, (num_head * num_tail, ))
topk_score, topk_sidx = th.topk(score,
k=topk if score.shape[0] > topk else score.shape[0],
dim=0)
sidx = th.argsort(topk_score, dim=0, descending=True)
score = topk_score[sidx]
sidx = topk_sidx[sidx]
idx = idx[sidx]
tail_idx = idx % num_tail
idx = floor_divide(idx, num_tail)
head_idx = idx % num_head
result.append((head[head_idx],
tail[tail_idx],
score))
else: # bcast at head
result = []
for i in range(num_head):
i_score = score[i]
topk_score, topk_sidx = th.topk(i_score,
k=topk if i_score.shape[0] > topk else i_score.shape[0],
dim=0)
sidx = th.argsort(topk_score, dim=0, descending=True)
i_score = topk_score[sidx]
idx = topk_sidx[sidx]
result.append((th.full((topk,), head[i], dtype=head[i].dtype),
tail[idx],
i_score))
return result