def forward_test_wikikg(self, query, ans, candidate, mode, gpu_id=-1):
"""Do the forward and generate ranking results.
Parameters
----------
pos_g : DGLGraph
Graph holding positive edges.
neg_g : DGLGraph
Graph holding negative edges.
logs : List
Where to put results in.
gpu_id : int
Which gpu to accelerate the calculation. if -1 is provided, cpu is used.
"""
scores = self.predict_score_wikikg(query, candidate, mode, to_device=cuda, gpu_id=gpu_id, trace=False)
argsort = F.argsort(scores, dim=1, descending=True)
return argsort[:,:10]
def forward_test_wikikg(self,
query,
ans,
candidate,
mode,
logs,
gpu_id=-1):
"""Do the forward and generate ranking results.
Parameters
----------
query : Tensor
input head and relation for test or valid
ans : Tenseor
the correct tail entity index
cadidate : Tensor
negative sampled tail entity
"""
scores = self.predict_score_wikikg(query,
candidate,
mode,
to_device=cuda,
gpu_id=gpu_id,
trace=False)
if mode == "Valid":
batch_size = query.shape[0]
neg_scores = reshape(scores, batch_size, -1)
for i in range(batch_size):
ranking = F.asnumpy(
F.sum(neg_scores[i] >= neg_scores[i][ans[i]], dim=0) + 1)
logs.append({
'MRR': 1.0 / ranking,
'MR': float(ranking),
'HITS@1': 1.0 if ranking <= 1 else 0.0,
'HITS@3': 1.0 if ranking <= 3 else 0.0,
'HITS@10': 1.0 if ranking <= 10 else 0.0
})
else:
argsort = F.argsort(scores, dim=1, descending=True)
logs.append(argsort[:, :10])
def topK(self, head=None, rel=None, tail=None, exec_mode='all', k=10):
if head is None:
head = F.arange(0, self.model.num_entity)
else:
head = F.tensor(head)
if rel is None:
rel = F.arange(0, self.model.num_rel)
else:
rel = F.tensor(rel)
if tail is None:
tail = F.arange(0, self.model.num_entity)
else:
tail = F.tensor(tail)
num_head = F.shape(head)[0]
num_rel = F.shape(rel)[0]
num_tail = F.shape(tail)[0]
if exec_mode == 'triplet_wise':
result = []
assert num_head == num_rel, \
'For triplet wise exection mode, head, relation and tail lists should have same length'
assert num_head == num_tail, \
'For triplet wise exection mode, head, relation and tail lists should have same length'
raw_score = self.model.score(head, rel, tail, triplet_wise=True)
score = self.score_func(raw_score)
idx = F.arange(0, num_head)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
idx = idx[sidx]
result.append((F.asnumpy(head[idx]),
F.asnumpy(rel[idx]),
F.asnumpy(tail[idx]),
F.asnumpy(score)))
elif exec_mode == 'all':
result = []
raw_score = self.model.score(head, rel, tail)
score = self.score_func(raw_score)
idx = F.arange(0, num_head * num_rel * num_tail)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
idx = idx[sidx]
tail_idx = idx % num_tail
idx = floor_divide(idx, num_tail)
rel_idx = idx % num_rel
idx = floor_divide(idx, num_rel)
head_idx = idx % num_head
result.append((F.asnumpy(head[head_idx]),
F.asnumpy(rel[rel_idx]),
F.asnumpy(tail[tail_idx]),
F.asnumpy(score)))
elif exec_mode == 'batch_head':
result = []
for i in range(num_head):
raw_score = self.model.score(F.unsqueeze(head[i], 0), rel, tail)
score = self.score_func(raw_score)
idx = F.arange(0, num_rel * num_tail)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
idx = idx[sidx]
tail_idx = idx % num_tail
idx = floor_divide(idx, num_tail)
rel_idx = idx % num_rel
result.append((np.full((k,), F.asnumpy(head[i])),
F.asnumpy(rel[rel_idx]),
F.asnumpy(tail[tail_idx]),
F.asnumpy(score)))
elif exec_mode == 'batch_rel':
result = []
for i in range(num_rel):
raw_score = self.model.score(head, F.unsqueeze(rel[i], 0), tail)
score = self.score_func(raw_score)
idx = F.arange(0, num_head * num_tail)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
idx = idx[sidx]
tail_idx = idx % num_tail
idx = floor_divide(idx, num_tail)
head_idx = idx % num_head
result.append((F.asnumpy(head[head_idx]),
np.full((k,), F.asnumpy(rel[i])),
F.asnumpy(tail[tail_idx]),
F.asnumpy(score)))
elif exec_mode == 'batch_tail':
result = []
for i in range(num_tail):
raw_score = self.model.score(head, rel, F.unsqueeze(tail[i], 0))
score = self.score_func(raw_score)
idx = F.arange(0, num_head * num_rel)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
idx = idx[sidx]
rel_idx = idx % num_rel
idx = floor_divide(idx, num_rel)
head_idx = idx % num_head
result.append((F.asnumpy(head[head_idx]),
F.asnumpy(rel[rel_idx]),
np.full((k,), F.asnumpy(tail[i])),
F.asnumpy(score)))
else:
assert False, 'unknow execution mode type {}'.format(exec_mode)
return result
def topK(self, head=None, tail=None, bcast=False, pair_ws=False, k=10):
if head is None:
head = F.arange(0, self.emb.shape[0])
else:
head = F.tensor(head)
if tail is None:
tail = F.arange(0, self.emb.shape[0])
else:
tail = F.tensor(tail)
head_emb = self.emb[head]
tail_emb = self.emb[tail]
if pair_ws is True:
result = []
batch_size = self.batch_size
# chunked cal score
score = []
num_head = head.shape[0]
num_tail = tail.shape[0]
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 = F.copy_to(sh_emb, 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 = F.copy_to(st_emb, self.device)
score.append(F.copy_to(self.sim_func(sh_emb, st_emb, pw=True), F.cpu()))
score = F.cat(score, dim=0)
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
result.append((F.asnumpy(head[sidx]),
F.asnumpy(tail[sidx]),
F.asnumpy(score)))
else:
num_head = head.shape[0]
num_tail = tail.shape[0]
batch_size = self.batch_size
# chunked cal score
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 = F.copy_to(sh_emb, 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 = F.copy_to(st_emb, self.device)
s_score.append(F.copy_to(self.sim_func(sh_emb, st_emb), F.cpu()))
score.append(F.cat(s_score, dim=1))
score = F.cat(score, dim=0)
if bcast is False:
result = []
idx = F.arange(0, num_head * num_tail)
score = F.reshape(score, (num_head * num_tail, ))
sidx = F.argsort(score, dim=0, descending=True)
sidx = sidx[:k]
score = score[sidx]
sidx = sidx
idx = idx[sidx]
tail_idx = idx % num_tail
idx = floor_divide(idx, num_tail)
head_idx = idx % num_head
result.append((F.asnumpy(head[head_idx]),
F.asnumpy(tail[tail_idx]),
F.asnumpy(score)))
else: # bcast at head
result = []
for i in range(num_head):
i_score = score[i]
sidx = F.argsort(i_score, dim=0, descending=True)
idx = F.arange(0, num_tail)
i_idx = sidx[:k]
i_score = i_score[i_idx]
idx = idx[i_idx]
result.append((np.full((k,), F.asnumpy(head[i])),
F.asnumpy(tail[idx]),
F.asnumpy(i_score)))
return result