def get_per_graph_ent_dropout_embeds(self, cur_time_list, target_time_list, node_sizes, time_diff_tensor, first_prev_graph_embeds, second_prev_graph_embeds):
batched_graph = self.get_batch_graph_dropout_embeds(filter_none(cur_time_list), target_time_list)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_embeds, second_layer_embeds = self.ent_encoder(batched_graph, first_prev_graph_embeds, second_prev_graph_embeds, time_diff_tensor, cur_time_list, node_sizes)
return first_layer_embeds.split(node_sizes), second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self, t_list, graph_train_list, val=False):
if val:
sampled_graph_list = graph_train_list
else:
sampled_graph_list = []
for g in graph_train_list:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
half_num_nodes = int(src.shape[0] / 2)
graph_split_ids = np.random.choice(np.arange(half_num_nodes),
size=int(0.5 * half_num_nodes), replace=False)
graph_split_rev_ids = graph_split_ids + half_num_nodes
sg = g.edge_subgraph(np.concatenate((graph_split_ids, graph_split_rev_ids)), preserve_nodes=True)
norm = comp_deg_norm(sg)
sg.ndata.update({'id': g.ndata['id'], 'norm': torch.from_numpy(norm).view(-1, 1)})
sg.edata['type_s'] = rel[np.concatenate((graph_split_ids, graph_split_rev_ids))]
sg.ids = g.ids
sampled_graph_list.append(sg)
time_embeds = []
for t, g in zip(t_list, graph_train_list):
temp_ent_embeds = torch.sin(t * self.w_ent_embeds[g.ndata['id']].view(-1, self.embed_size) +
self.b_ent_embeds[g.ndata['id']].view(-1, self.embed_size))
time_embeds.append(temp_ent_embeds)
batched_graph = dgl.batch(sampled_graph_list)
time_embeds = torch.cat(time_embeds, dim=0)
ent_embeds = self.ent_embeds[batched_graph.ndata['id']].view(-1, self.embed_size)
batched_graph.ndata['h'] = torch.cat([ent_embeds, time_embeds], dim=-1)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
node_sizes = [len(g.nodes()) for g in graph_train_list]
enc_ent_mean_graph = self.ent_encoder(batched_graph, reverse=False)
ent_enc_embeds = enc_ent_mean_graph.ndata['h']
per_graph_ent_embeds = ent_enc_embeds.split(node_sizes)
return per_graph_ent_embeds
def get_per_graph_ent_embeds(self, g_batched_list_t, time_batched_list_t, node_sizes, time_diff_tensor, first_prev_graph_embeds, second_prev_graph_embeds, full, rate=0.5):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t, full, rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_embeds, second_layer_embeds = self.ent_encoder(batched_graph, first_prev_graph_embeds, second_prev_graph_embeds, time_diff_tensor, time_batched_list_t, node_sizes)
return first_layer_embeds.split(node_sizes), second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self,
g_batched_list_t,
time_batched_list_t,
node_sizes,
full=False,
rate=0.5):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t, full,
rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_embeds, second_layer_embeds = self.ent_encoder(
batched_graph, time_batched_list_t, node_sizes)
# first_layer_embeds = first_layer_graph.ndata['h']
# second_layer_embeds = second_layer_graph.ndata['h']
'''
first_layer_embeds_split = []
second_layer_embeds_split = []
for t, first_layer_embed, second_layer_embed in zip(time_batched_list_t, first_layer_embeds.split(node_sizes), second_layer_embeds.split(node_sizes)):
first_layer_embeds_split.append(first_layer_embed + self.time_embed[t])
second_layer_embeds_split.append(second_layer_embed + self.time_embed[t])
return first_layer_embeds_split, second_layer_embeds_split
'''
return first_layer_embeds.split(node_sizes), second_layer_embeds.split(
node_sizes)
def get_graph_embeds_center(self,
g_batched_list_t,
time_batched_list_t,
node_sizes,
first_prev_graph_embeds_forward,
second_prev_graph_embeds_forward,
time_diff_tensor_forward,
first_prev_graph_embeds_backward,
second_prev_graph_embeds_backward,
time_diff_tensor_backward,
full,
rate=0.5):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t, full,
rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
second_local_embeds, second_layer_embeds = self.ent_encoder.forward_post_ensemble(
batched_graph, first_prev_graph_embeds_forward,
second_prev_graph_embeds_forward, time_diff_tensor_forward,
first_prev_graph_embeds_backward,
second_prev_graph_embeds_backward, time_diff_tensor_backward,
time_batched_list_t, node_sizes)
return second_local_embeds.split(
node_sizes), second_layer_embeds.split(node_sizes)
def get_final_graph_embeds(self, g_batched_list_t, time_batched_list_t, node_sizes, hist_embeddings, attn_mask, full, rate=0.5, val=False):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t, full, rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_prev_embeddings, second_layer_prev_embeddings, local_attn_mask = self.get_prev_embeddings(g_batched_list_t, hist_embeddings, attn_mask)
second_local_embeds, second_layer_embeds = self.ent_encoder.forward_post_ensemble(batched_graph, second_layer_prev_embeddings,
self.time_diff_test if val else self.time_diff_train, local_attn_mask, time_batched_list_t, node_sizes)
return second_local_embeds.split(node_sizes), second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self, t_list, graph_train_list, val=False):
if val:
sampled_graph_list = graph_train_list
else:
sampled_graph_list = []
for g in graph_train_list:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
total_idx = np.random.choice(np.arange(src.shape[0]),
size=int(0.5 * src.shape[0]),
replace=False)
sg = g.edge_subgraph(total_idx, preserve_nodes=True)
node_norm = comp_deg_norm(sg)
sg.ndata.update({
'id': g.ndata['id'],
'norm': torch.from_numpy(node_norm).view(-1, 1)
})
sg.edata['norm'] = node_norm_to_edge_norm(
sg,
torch.from_numpy(node_norm).view(-1, 1))
sg.edata['type_s'] = rel[total_idx]
sg.ids = g.ids
sampled_graph_list.append(sg)
# time_embeds = []
# for t, g in zip(t_list, graph_train_list):
# temp_ent_embeds = torch.sin(t * self.w_ent_embeds[g.ndata['id']].view(-1, self.embed_size) +
# self.b_ent_embeds[g.ndata['id']].view(-1, self.embed_size))
# time_embeds.append(temp_ent_embeds)
ent_embeds = []
for t, g in zip(t_list, graph_train_list):
static_ent_embeds = self.ent_embeds[g.ndata['id']].view(
-1, self.embed_size)
ones = static_ent_embeds.new_ones(static_ent_embeds.shape[0],
self.static_embed_size)
temp_ent_embeds = torch.sin(t * self.w_temp_ent_embeds[
g.ndata['id']].view(-1, self.temporal_embed_size) +
self.b_temp_ent_embeds[g.ndata['id']].
view(-1, self.temporal_embed_size))
ent_embeds.append(static_ent_embeds * torch.cat(
(ones, temp_ent_embeds), dim=-1))
batched_graph = dgl.batch(sampled_graph_list)
batched_graph.ndata['h'] = torch.cat(ent_embeds, dim=0)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
node_sizes = [len(g.nodes()) for g in graph_train_list]
enc_ent_mean_graph = self.ent_encoder(batched_graph)
ent_enc_embeds = enc_ent_mean_graph.ndata['h']
per_graph_ent_embeds = ent_enc_embeds.split(node_sizes)
return per_graph_ent_embeds
def get_per_graph_ent_dropout_embeds_one_direction(
self, cur_time_list, target_time_list, node_sizes,
time_diff_tensor, first_prev_graph_embeds,
second_prev_graph_embeds, forward):
batched_graph = self.get_batch_graph_dropout_embeds(
filter_none(cur_time_list), target_time_list)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
second_local_embeds, first_layer_embeds, second_layer_embeds = self.ent_encoder.forward_post_ensemble_one_direction(
batched_graph, first_prev_graph_embeds, second_prev_graph_embeds,
time_diff_tensor, cur_time_list, node_sizes, forward)
return second_local_embeds.split(node_sizes), first_layer_embeds.split(
node_sizes), second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self,
g_batched_list_t,
node_sizes,
first_prev_graph_embeds,
second_prev_graph_embeds,
val=False):
if val:
sampled_graph_list = g_batched_list_t
else:
sampled_graph_list = []
for g in g_batched_list_t:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
half_num_nodes = int(src.shape[0] / 2)
graph_split_ids = np.random.choice(np.arange(half_num_nodes),
size=int(0.5 *
half_num_nodes),
replace=False)
graph_split_rev_ids = graph_split_ids + half_num_nodes
sg = g.edge_subgraph(np.concatenate(
(graph_split_ids, graph_split_rev_ids)),
preserve_nodes=True)
node_norm = comp_deg_norm(sg)
sg.ndata.update({
'id': g.ndata['id'],
'norm': torch.from_numpy(node_norm).view(-1, 1)
})
sg.edata['norm'] = node_norm_to_edge_norm(
sg,
torch.from_numpy(node_norm).view(-1, 1))
sg.edata['type_s'] = rel[np.concatenate(
(graph_split_ids, graph_split_rev_ids))]
sg.ids = g.ids
sampled_graph_list.append(sg)
batched_graph = dgl.batch(sampled_graph_list)
batched_graph.ndata['h'] = self.ent_embeds[
batched_graph.ndata['id']].view(-1, self.embed_size)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_graph, second_layer_graph = self.ent_encoder(
batched_graph, first_prev_graph_embeds, second_prev_graph_embeds)
first_layer_embeds = first_layer_graph.ndata['h']
second_layer_embeds = second_layer_graph.ndata['h']
return first_layer_embeds.split(node_sizes), second_layer_embeds.split(
node_sizes)
def get_per_graph_ent_embeds(self,
g_batched_list_t,
cur_h,
node_sizes,
val=False):
if val:
sampled_graph_list = g_batched_list_t
else:
sampled_graph_list = []
for g in g_batched_list_t:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
half_num_nodes = int(src.shape[0] / 2)
graph_split_ids = np.random.choice(np.arange(half_num_nodes),
size=int(0.5 *
half_num_nodes),
replace=False)
graph_split_rev_ids = graph_split_ids + half_num_nodes
sg = g.edge_subgraph(np.concatenate(
(graph_split_ids, graph_split_rev_ids)),
preserve_nodes=True)
norm = comp_deg_norm(sg)
sg.ndata.update({
'id': g.ndata['id'],
'norm': torch.from_numpy(norm).view(-1, 1)
})
sg.edata['type_s'] = rel[np.concatenate(
(graph_split_ids, graph_split_rev_ids))]
sg.ids = g.ids
sampled_graph_list.append(sg)
batched_graph = dgl.batch(sampled_graph_list)
expanded_h = torch.cat([
cur_h[i].unsqueeze(0).expand(size, self.embed_size)
for i, size in enumerate(node_sizes)
],
dim=0)
ent_embeds = self.ent_embeds[batched_graph.ndata['id']].view(
-1, self.embed_size)
batched_graph.ndata['h'] = torch.cat([ent_embeds, expanded_h], dim=-1)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
enc_ent_mean_graph = self.ent_encoder(batched_graph, reverse=False)
ent_enc_embeds = enc_ent_mean_graph.ndata['h']
per_graph_ent_embeds = ent_enc_embeds.split(node_sizes)
return per_graph_ent_embeds
def get_per_graph_ent_embeds_one_direction(self,
g_batched_list_t,
time_batched_list_t,
node_sizes,
time_diff_tensor,
first_prev_graph_embeds,
second_prev_graph_embeds,
forward,
full,
rate=0.5):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t,
full=full,
rate=rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
second_local_embeds, first_layer_embeds, second_layer_embeds = self.ent_encoder.forward_post_ensemble_one_direction(
batched_graph, first_prev_graph_embeds, second_prev_graph_embeds,
time_diff_tensor, time_batched_list_t, node_sizes, forward)
return second_local_embeds.split(node_sizes), first_layer_embeds.split(
node_sizes), second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self, t_list, graph_train_list, val=False):
if val:
sampled_graph_list = graph_train_list
else:
# TODO: modify half_num_nodes
sampled_graph_list = []
for g in graph_train_list:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
half_num_nodes = int(src.shape[0] / 2)
# graph_split_ids = np.random.choice(np.arange(half_num_nodes),
# size=int(0.5 * half_num_nodes), replace=False)
# graph_split_rev_ids = graph_split_ids + half_num_nodes
# sg = g.edge_subgraph(np.concatenate((graph_split_ids, graph_split_rev_ids)), preserve_nodes=True)
total_idx = np.random.choice(np.arange(src.shape[0]),
size=int(0.5 * src.shape[0]),
replace=False)
sg = g.edge_subgraph(total_idx, preserve_nodes=True)
node_norm = comp_deg_norm(sg)
sg.ndata.update({
'id': g.ndata['id'],
'norm': torch.from_numpy(node_norm).view(-1, 1)
})
sg.edata['norm'] = node_norm_to_edge_norm(
sg,
torch.from_numpy(node_norm).view(-1, 1))
sg.edata['type_s'] = rel[total_idx]
sg.ids = g.ids
sampled_graph_list.append(sg)
batched_graph = dgl.batch(sampled_graph_list)
batched_graph.ndata['h'] = self.ent_embeds[
batched_graph.ndata['id']].view(-1, self.embed_size)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
node_sizes = [len(g.nodes()) for g in graph_train_list]
enc_ent_mean_graph = self.ent_encoder(batched_graph, t_list,
node_sizes)
ent_enc_embeds = enc_ent_mean_graph.ndata['h']
per_graph_ent_embeds = ent_enc_embeds.split(node_sizes)
return per_graph_ent_embeds
def get_final_graph_embeds(self,
g_batched_list_t,
time_batched_list_t,
node_sizes,
hist_embeddings,
attn_mask,
full,
rate=0.5,
val=False):
batched_graph = self.get_batch_graph_embeds(g_batched_list_t, full,
rate)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_prev_embeddings, second_layer_prev_embeddings, local_attn_mask = self.get_prev_embeddings(
g_batched_list_t, hist_embeddings, attn_mask)
# if self.EMA:
# second_layer_embeds = self.ent_encoder.forward_ema(batched_graph, second_layer_prev_embeddings, time_batched_list_t, node_sizes, torch.sigmoid(self.alpha), self.train_seq_len)
second_layer_embeds = self.ent_encoder.forward_final(
batched_graph, first_layer_prev_embeddings,
second_layer_prev_embeddings,
self.time_diff_test if val else self.time_diff_train,
local_attn_mask, time_batched_list_t, node_sizes)
return second_layer_embeds.split(node_sizes)
def get_per_graph_ent_embeds(self,
g_batched_list_t,
time_batched_list_t,
node_sizes,
time_diff_tensor,
first_prev_graph_embeds,
second_prev_graph_embeds,
val=False):
if val:
sampled_graph_list = g_batched_list_t
else:
sampled_graph_list = []
for g in g_batched_list_t:
src, rel, dst = g.edges()[0], g.edata['type_s'], g.edges()[1]
half_num_nodes = int(src.shape[0] / 2)
# graph_split_ids = np.random.choice(np.arange(half_num_nodes), size=int(0.5 * src.shape[0]), replace=False)
# graph_split_rev_ids = graph_split_ids + half_num_nodes
# total_idx = np.concatenate((graph_split_ids, graph_split_rev_ids))
total_idx = np.random.choice(np.arange(src.shape[0]),
size=int(0.5 * src.shape[0]),
replace=False)
sg = g.edge_subgraph(total_idx, preserve_nodes=True)
node_norm = comp_deg_norm(sg)
sg.ndata.update({
'id': g.ndata['id'],
'norm': torch.from_numpy(node_norm).view(-1, 1)
})
sg.edata['norm'] = node_norm_to_edge_norm(
sg,
torch.from_numpy(node_norm).view(-1, 1))
sg.edata['type_s'] = rel[total_idx]
sg.ids = g.ids
sampled_graph_list.append(sg)
ent_embeds = []
for t, g in zip(time_batched_list_t, g_batched_list_t):
static_ent_embeds = self.ent_embeds[g.ndata['id']].view(
-1, self.embed_size)
ones = static_ent_embeds.new_ones(static_ent_embeds.shape[0],
self.static_embed_size)
temp_ent_embeds = torch.sin(t * self.w_temp_ent_embeds[
g.ndata['id']].view(-1, self.temporal_embed_size) +
self.b_temp_ent_embeds[g.ndata['id']].
view(-1, self.temporal_embed_size))
ent_embeds.append(static_ent_embeds * torch.cat(
(ones, temp_ent_embeds), dim=-1))
batched_graph = dgl.batch(sampled_graph_list)
batched_graph.ndata['h'] = torch.cat(ent_embeds, dim=0)
if self.use_cuda:
move_dgl_to_cuda(batched_graph)
first_layer_graph, second_layer_graph = self.ent_encoder(
batched_graph, first_prev_graph_embeds, second_prev_graph_embeds,
time_diff_tensor)
first_layer_embeds = first_layer_graph.ndata['h']
second_layer_embeds = second_layer_graph.ndata['h']
return first_layer_embeds.split(node_sizes), second_layer_embeds.split(
node_sizes)
