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)