文章目錄
- @[toc]
- TGN模型訓練階段代碼理解
- 論文信息
- 代碼過程手繪
- 代碼訓練過程
- compute_temporal_embeddings
- update_memory
- get_raw_messages
- get_updated_memory
- self.message_aggregator.aggregate
- self.memory_updater.get_updated_memory
- Memory
- get_embedding_module
- GraphAttentionEmbedding
- TimeEncode
- NeighborFinder
- MergeLayer
文章目錄
- @[toc]
- TGN模型訓練階段代碼理解
- 論文信息
- 代碼過程手繪
- 代碼訓練過程
- compute_temporal_embeddings
- update_memory
- get_raw_messages
- get_updated_memory
- self.message_aggregator.aggregate
- self.memory_updater.get_updated_memory
- Memory
- get_embedding_module
- GraphAttentionEmbedding
- TimeEncode
- NeighborFinder
- MergeLayer
TGN模型訓練階段代碼理解
論文信息
論文鏈接：https://arxiv.org/abs/2006.10637
GitHub: https://github.com/twitter-research/tgn?tab=readme-ov-file
年份：2020
代碼過程手繪
代碼訓練過程
pos_prob, neg_prob = tgn.compute_edge_probabilities(sources_batch, destinations_batch, negatives_batch,timestamps_batch, edge_idxs_batch, NUM_NEIGHBORS)
函數compute_edge_probabilities
def compute_edge_probabilities(self, source_nodes, destination_nodes, negative_nodes, edge_times,edge_idxs, n_neighbors=20):"""Compute probabilities for edges between sources and destination and between sources andnegatives by first computing temporal embeddings using the TGN encoder and then feeding theminto the MLP decoder.:param destination_nodes [batch_size]: destination ids:param negative_nodes [batch_size]: ids of negative sampled destination:param edge_times [batch_size]: timestamp of interaction:param edge_idxs [batch_size]: index of interaction:param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutionallayer:return: Probabilities for both the positive and negative edgessource_nodes 源節點id列表destination_nodes 目標節點id列表negative_nodes 負采樣節點id列表edge_times 源節點列表中的節點與目標節點列表中的節點發生關系時的時間edge_idxs 邊的編號"""n_samples = len(source_nodes)# compute_temporal_embeddingssource_node_embedding, destination_node_embedding, negative_node_embedding = self.compute_temporal_embeddings(source_nodes, destination_nodes, negative_nodes, edge_times, edge_idxs, n_neighbors)score = self.affinity_score(torch.cat([source_node_embedding, source_node_embedding], dim=0),torch.cat([destination_node_embedding,negative_node_embedding])).squeeze(dim=0)pos_score = score[:n_samples]neg_score = score[n_samples:]return pos_score.sigmoid(), neg_score.sigmoid()
compute_temporal_embeddings
這個方法的目的是計算時間嵌入
def compute_temporal_embeddings(self, source_nodes, destination_nodes, negative_nodes, edge_times,edge_idxs, n_neighbors=20):"""Compute temporal embeddings for sources, destinations, and negatively sampled destinations.這個方法的目的是計算時間嵌入source_nodes [batch_size]: source ids.:param destination_nodes [batch_size]: destination ids:param negative_nodes [batch_size]: ids of negative sampled destination:param edge_times [batch_size]: timestamp of interaction:param edge_idxs [batch_size]: index of interaction:param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutionallayer:return: Temporal embeddings for sources, destinations and negatives"""# n_samples 表示源節點有多少個n_samples = len(source_nodes)# nodes是所有的節點這個batch_size中所有的節點id, size=200*3=600nodes = np.concatenate([source_nodes, destination_nodes, negative_nodes])# positives 是將源節點和目標節點和在一切，前200個是源節點的node_id, 后200個是目標節點的node_idpositives = np.concatenate([source_nodes, destination_nodes])# timestamps shape=200*3 edge_times 是發生交互的時間timestamps = np.concatenate([edge_times, edge_times, edge_times])# edge_times shape = batch_size 是源節點和目的節點發生的時間memory = Nonetime_diffs = Noneif self.use_memory:if self.memory_update_at_start: # 是不是剛開始使用記憶# n_nodes 表示的是圖中一共有多少個節點 9228# 記憶列表 self.memory.messages 當前狀態一定為空# 在這個地方出來的memory是最新的memory，是根據節點的messages信息進行更新的，在代碼中會取該節點messages列表中最新的那一個memory, last_update = self.get_updated_memory(list(range(self.n_nodes)),self.memory.messages) # memory shape = [n_nodes(9228), memory_dimension(172)] last_update shape [n_nodes(9228)]else:memory = self.memory.get_memory(list(range(self.n_nodes)))last_update = self.memory.last_update# ===================================== 下面這些都是處理單個節點的信息 ==============================# 計算節點內存最后一次更新的時間與我們希望計算該節點嵌入的時間之間的差異。# source_time_diffs shape [batch_size]source_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[source_nodes].long() # 這是標準化操作source_time_diffs = (source_time_diffs - self.mean_time_shift_src) / self.std_time_shift_src # destination_time_diffs shape [batch_size]destination_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[destination_nodes].long()# 這是標準化操作destination_time_diffs = (destination_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dstnegative_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[negative_nodes].long()negative_time_diffs = (negative_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dst# 時間差time_diffs = torch.cat([source_time_diffs, destination_time_diffs, negative_time_diffs],dim=0)# Compute the embeddings using the embedding module# self.embedding_module 在下面所示# 1. 先是 self.embedding_module = get_embedding_module()"""memory 記憶對象nodes 是一個結合了源節點目的節點和負采樣節點的node_id列表timestamps 200*3的時間列表self.n_layers 遞歸的層數 這里為2n_neighbors 選取多少個鄰居節點 這里是10time_diffs 標準化過后的時間差"""# node_embedding shape [600, 172] 融合了節點的特征和鄰居其余邊的特征node_embedding = self.embedding_module.compute_embedding(memory=memory,source_nodes=nodes,timestamps=timestamps,n_layers=self.n_layers,n_neighbors=n_neighbors,time_diffs=time_diffs)# 然后去獲取不同列表的節點特征 source_node_embedding = node_embedding[:n_samples]destination_node_embedding = node_embedding[n_samples: 2 * n_samples]negative_node_embedding = node_embedding[2 * n_samples:]if self.use_memory:# 進行記憶力更新if self.memory_update_at_start:# Persist the updates to the memory only for sources and destinations (since now we have# new messages for them)self.update_memory(positives, self.memory.messages)assert torch.allclose(memory[positives], self.memory.get_memory(positives), atol=1e-5), \"Something wrong in how the memory was updated"# Remove messages for the positives since we have already updated the memory using them# 記憶已經更新了，那么對于每個信息就即為空self.memory.clear_messages(positives)unique_sources, source_id_to_messages = self.get_raw_messages(source_nodes,source_node_embedding,destination_nodes,destination_node_embedding,edge_times, edge_idxs)unique_destinations, destination_id_to_messages = self.get_raw_messages(destination_nodes,destination_node_embedding,source_nodes,source_node_embedding,edge_times, edge_idxs)if self.memory_update_at_start:# 存儲信息self.memory.store_raw_messages(unique_sources, source_id_to_messages)self.memory.store_raw_messages(unique_destinations, destination_id_to_messages)else:self.update_memory(unique_sources, source_id_to_messages)self.update_memory(unique_destinations, destination_id_to_messages)if self.dyrep:source_node_embedding = memory[source_nodes]destination_node_embedding = memory[destination_nodes]negative_node_embedding = memory[negative_nodes]return source_node_embedding, destination_node_embedding, negative_node_embedding
update_memory
def update_memory(self, nodes, messages):# Aggregate messages for the same nodes# self.message_aggregator -> LastMessageAggregatorunique_nodes, unique_messages, unique_timestamps = \self.message_aggregator.aggregate(nodes,messages)if len(unique_nodes) > 0:unique_messages = self.message_function.compute_message(unique_messages)# Update the memory with the aggregated messages# 聚合完了就去更新self.memory_updater.update_memory(unique_nodes, unique_messages,timestamps=unique_timestamps)
get_raw_messages
def get_raw_messages(self, source_nodes, source_node_embedding, destination_nodes,destination_node_embedding, edge_times, edge_idxs):# edge_times shape is [200, ]edge_times = torch.from_numpy(edge_times).float().to(self.device)# edge_features shape is [200, 172]edge_features = self.edge_raw_features[edge_idxs]source_memory = self.memory.get_memory(source_nodes) if not \self.use_source_embedding_in_message else source_node_embeddingdestination_memory = self.memory.get_memory(destination_nodes) if \not self.use_destination_embedding_in_message else destination_node_embeddingsource_time_delta = edge_times - self.memory.last_update[source_nodes]# source_time_delta_encoding [200, 172]source_time_delta_encoding = self.time_encoder(source_time_delta.unsqueeze(dim=1)).view(len(source_nodes), -1)# source_message shape [200, 688]source_message = torch.cat([source_memory, destination_memory, edge_features,source_time_delta_encoding],dim=1)messages = defaultdict(list)unique_sources = np.unique(source_nodes)for i in range(len(source_nodes)):messages[source_nodes[i]].append((source_message[i], edge_times[i]))return unique_sources, messages
get_updated_memory
def get_updated_memory(self, nodes, messages):# Aggregate messages for the same nodes# nodes 是一個列表 range(n_nodes)# messages是消息列表# 先是聚合消息，然后更新記憶# 在第一次進來這個函數的時候，返回的全是[]unique_nodes, unique_messages, unique_timestamps = \self.message_aggregator.aggregate(nodes, # 是一個列表 range(n_nodes)messages # 是消息列表)if len(unique_nodes) > 0:# 有兩個選擇"""class MLPMessageFunction(MessageFunction):def __init__(self, raw_message_dimension, message_dimension):super(MLPMessageFunction, self).__init__()self.mlp = self.layers = nn.Sequential(nn.Linear(raw_message_dimension, raw_message_dimension // 2),nn.ReLU(),nn.Linear(raw_message_dimension // 2, message_dimension),)def compute_message(self, raw_messages):messages = self.mlp(raw_messages)return messagesclass IdentityMessageFunction(MessageFunction): def compute_message(self, raw_messages):# 作者使用的是這個，啥也沒有邊，直接返回return raw_messages"""unique_messages = self.message_function.compute_message(unique_messages)# 在頭一次訓練的過程中進來這個地方, 返回的全是0的矩陣# 形狀為，[n_nodes, memory_dimension] [n_nodes]updated_memory, updated_last_update = self.memory_updater.get_updated_memory(unique_nodes,unique_messages,timestamps=unique_timestamps)return updated_memory, updated_last_update
self.message_aggregator.aggregate
代碼中默認使用last
def get_message_aggregator(aggregator_type, device):if aggregator_type == "last":return LastMessageAggregator(device=device)elif aggregator_type == "mean":return MeanMessageAggregator(device=device)else:raise ValueError("Message aggregator {} not implemented".format(aggregator_type))
LastMessageAggregator代碼：
class LastMessageAggregator(MessageAggregator):def __init__(self, device):super(LastMessageAggregator, self).__init__(device)def aggregate(self, node_ids, messages):"""Only keep the last message for each node"""unique_node_ids = np.unique(node_ids) # 去重節點，不知道啥作用，因為本來就沒有重復unique_messages = []unique_timestamps = []to_update_node_ids = []for node_id in unique_node_ids: # 循環range(n_nodes)=9228if len(messages[node_id]) > 0:"""上一步結束每個節點存儲的信息以及對應的（時間？）source_message = torch.cat([source_memory, destination_memory, edge_features,source_time_delta_encoding], dim=1)source_message, edge_times"""to_update_node_ids.append(node_id)unique_messages.append(messages[node_id][-1][0])unique_timestamps.append(messages[node_id][-1][1])unique_messages = torch.stack(unique_messages) if len(to_update_node_ids) > 0 else []unique_timestamps = torch.stack(unique_timestamps) if len(to_update_node_ids) > 0 else []return to_update_node_ids, unique_messages, unique_timestamps
self.memory_updater.get_updated_memory
代碼中默認采用使用gru的方式去更新記憶
class SequenceMemoryUpdater(MemoryUpdater):def __init__(self, memory, message_dimension, memory_dimension, device):super(SequenceMemoryUpdater, self).__init__()self.memory = memoryself.layer_norm = torch.nn.LayerNorm(memory_dimension)self.message_dimension = message_dimensionself.device = devicedef update_memory(self, unique_node_ids, unique_messages, timestamps):if len(unique_node_ids) <= 0:returnassert (self.memory.get_last_update(unique_node_ids) <= timestamps).all().item(), "Trying to " \"update memory to time in the past"memory = self.memory.get_memory(unique_node_ids)self.memory.last_update[unique_node_ids] = timestampsupdated_memory = self.memory_updater(unique_messages, memory)self.memory.set_memory(unique_node_ids, updated_memory)def get_updated_memory(self, unique_node_ids, unique_messages, timestamps):if len(unique_node_ids) <= 0:# 這里的self.memory在下面進行定義# self.memory.memory 在初始化的時候是一個全為0，shape=[n_nodes, memory_dimension], 沒有梯度的矩陣# self.memory.last_update 在初始化的時候是一個全為0，shape=[n_nodes], 沒有梯度的舉證# 這里的clone是深拷貝，并不會影響原來的值是多少# 第二次就不是走這里咯return self.memory.memory.data.clone(), self.memory.last_update.data.clone()assert (self.memory.get_last_update(unique_node_ids) <= timestamps).all().item(), "Trying to " \"update memory to time in the past"updated_memory = self.memory.memory.data.clone()updated_memory[unique_node_ids] = self.memory_updater(unique_messages, updated_memory[unique_node_ids])updated_last_update = self.memory.last_update.data.clone()updated_last_update[unique_node_ids] = timestampsreturn updated_memory, updated_last_updateclass GRUMemoryUpdater(SequenceMemoryUpdater):def __init__(self, memory, message_dimension, memory_dimension, device):super(GRUMemoryUpdater, self).__init__(memory, message_dimension, memory_dimension, device)self.memory_updater = nn.GRUCell(input_size=message_dimension,hidden_size=memory_dimension)
Memory
class Memory(nn.Module):def __init__(self, n_nodes, memory_dimension, input_dimension, message_dimension=None,device="cpu", combination_method='sum'):super(Memory, self).__init__()self.n_nodes = n_nodesself.memory_dimension = memory_dimensionself.input_dimension = input_dimensionself.message_dimension = message_dimensionself.device = deviceself.combination_method = combination_methodself.__init_memory__()# 這是初是化def __init_memory__(self):"""Initializes the memory to all zeros. It should be called at the start of each epoch."""# Treat memory as parameter so that it is saved and loaded together with the model# self.memory_dimension = 172# self.n_nodes = 9228# self.memory shape is [9228, 172]的一個記憶，每一個節點都有對應的記憶，并且每一個記憶向量是172# self.memory = 一個全為0的矩陣self.memory = nn.Parameter(torch.zeros((self.n_nodes, self.memory_dimension)).to(self.device),requires_grad=False)# last_update shape = [9228]self.last_update = nn.Parameter(torch.zeros(self.n_nodes).to(self.device),requires_grad=False)self.messages = defaultdict(list)def store_raw_messages(self, nodes, node_id_to_messages):for node in nodes:self.messages[node].extend(node_id_to_messages[node])def get_memory(self, node_idxs):return self.memory[node_idxs, :]def set_memory(self, node_idxs, values):self.memory[node_idxs, :] = valuesdef get_last_update(self, node_idxs):return self.last_update[node_idxs]def backup_memory(self):messages_clone = {}for k, v in self.messages.items():messages_clone[k] = [(x[0].clone(), x[1].clone()) for x in v]return self.memory.data.clone(), self.last_update.data.clone(), messages_clonedef restore_memory(self, memory_backup):self.memory.data, self.last_update.data = memory_backup[0].clone(), memory_backup[1].clone()self.messages = defaultdict(list)for k, v in memory_backup[2].items():self.messages[k] = [(x[0].clone(), x[1].clone()) for x in v]def detach_memory(self):self.memory.detach_()# Detach all stored messagesfor k, v in self.messages.items():new_node_messages = []for message in v:new_node_messages.append((message[0].detach(), message[1]))self.messages[k] = new_node_messagesdef clear_messages(self, nodes):for node in nodes:self.messages[node] = []
get_embedding_module
這里的module_type=graph_attention
def get_embedding_module(module_type, node_features, edge_features, memory, neighbor_finder,time_encoder, n_layers, n_node_features, n_edge_features, n_time_features,embedding_dimension, device,n_heads=2, dropout=0.1, n_neighbors=None,use_memory=True):# embedding_module采用的是這個if module_type == "graph_attention":return GraphAttentionEmbedding(node_features=node_features,edge_features=edge_features,memory=memory,neighbor_finder=neighbor_finder,time_encoder=time_encoder,n_layers=n_layers,n_node_features=n_node_features,n_edge_features=n_edge_features,n_time_features=n_time_features,embedding_dimension=embedding_dimension,device=device,n_heads=n_heads, dropout=dropout, use_memory=use_memory)elif module_type == "graph_sum":return GraphSumEmbedding(node_features=node_features,edge_features=edge_features,memory=memory,neighbor_finder=neighbor_finder,time_encoder=time_encoder,n_layers=n_layers,n_node_features=n_node_features,n_edge_features=n_edge_features,n_time_features=n_time_features,embedding_dimension=embedding_dimension,device=device,n_heads=n_heads, dropout=dropout, use_memory=use_memory)elif module_type == "identity":return IdentityEmbedding(node_features=node_features,edge_features=edge_features,memory=memory,neighbor_finder=neighbor_finder,time_encoder=time_encoder,n_layers=n_layers,n_node_features=n_node_features,n_edge_features=n_edge_features,n_time_features=n_time_features,embedding_dimension=embedding_dimension,device=device,dropout=dropout)elif module_type == "time":return TimeEmbedding(node_features=node_features,edge_features=edge_features,memory=memory,neighbor_finder=neighbor_finder,time_encoder=time_encoder,n_layers=n_layers,n_node_features=n_node_features,n_edge_features=n_edge_features,n_time_features=n_time_features,embedding_dimension=embedding_dimension,device=device,dropout=dropout,n_neighbors=n_neighbors)else:raise ValueError("Embedding Module {} not supported".format(module_type))
GraphAttentionEmbedding
class GraphEmbedding(EmbeddingModule):def __init__(self, node_features, edge_features, memory, neighbor_finder, time_encoder, n_layers,n_node_features, n_edge_features, n_time_features, embedding_dimension, device,n_heads=2, dropout=0.1, use_memory=True):super(GraphEmbedding, self).__init__(node_features, edge_features, memory,neighbor_finder, time_encoder, n_layers,n_node_features, n_edge_features, n_time_features,embedding_dimension, device, dropout)self.use_memory = use_memoryself.device = devicedef compute_embedding(self, memory, source_nodes, timestamps, n_layers, n_neighbors=20, time_diffs=None,use_time_proj=True):"""Recursive implementation of curr_layers temporal graph attention layers.使用遞歸的方式來實現一系列時間圖注意力src_idx_l [batch_size]: users / items input ids.cut_time_l [batch_size]: scalar representing the instant of the time where we want to extract the user / item representation.curr_layers [scalar]: number of temporal convolutional layers to stack.num_neighbors [scalar]: number of temporal neighbor to consider in each convolutional layer.""""""memory 記憶對象source_nodes 是一個結合了源節點目的節點和負采樣節點的node_id列表（一開始是，后面不是）timestamps 200*3的時間列表self.n_layers 遞歸的層數 這里為2n_neighbors 選取多少個鄰居節點 這里是10time_diffs 標準化過后的時間差"""assert (n_layers >= 0)# source_nodes_torch shape = [n_nodes]source_nodes_torch = torch.from_numpy(source_nodes).long().to(self.device)# timestamps_torch shape = [3*200, 1]timestamps_torch = torch.unsqueeze(torch.from_numpy(timestamps).float().to(self.device), dim=1)# query node always has the start time -> time span == 0# 這里的time_encoder是一個模型,經過的是一個cos(linear(x))，在下面有對應的代碼# torch.zeros_like(timestamps_torch) 是一個全為0 shape = [3*200, 1]# source_nodes_time_embedding shape = [3*200, 1, 172]source_nodes_time_embedding = self.time_encoder(torch.zeros_like(timestamps_torch))# self.node_features是一個全為0的矩陣# self.node_features shape is [n_nodes, node_dim] = [9228, 172]# source_node_features 是所有節點的特征 shape is [600, 172]source_node_features = self.node_features[source_nodes_torch, :]if self.use_memory:# 將節點當前的特征 再加上記憶中節點的特征source_node_features = memory[source_nodes, :] + source_node_features# ====================================== 這下面執行了一個遞歸的操作 ==================================# n_layers = 1if n_layers == 0:return source_node_featureselse:# 再一次調用自己，返回的是節點的特征shape is [600, 172]source_node_conv_embeddings = self.compute_embedding(memory,source_nodes,timestamps,n_layers=n_layers - 1,n_neighbors=n_neighbors)# 獲得是source_nodes這3*200個節點，在3*200的時間列表中，選取前十個鄰居"""neighbors shape is [3*200, n_neighbors]edge_idxs shape is [3*200, n_neighbors]edge_times shape is [3*200, n_neighbors] """neighbors, edge_idxs, edge_times = self.neighbor_finder.get_temporal_neighbor(source_nodes,timestamps,n_neighbors=n_neighbors)# 這里的鄰居節點node_id是source_nodes中的每一個鄰居節點，變成torch形式neighbors_torch = torch.from_numpy(neighbors).long().to(self.device)edge_idxs = torch.from_numpy(edge_idxs).long().to(self.device)# 時間差，600個節點的edge_deltas = timestamps[:, np.newaxis] - edge_timesedge_deltas_torch = torch.from_numpy(edge_deltas).float().to(self.device)# 展平，變成6000neighbors = neighbors.flatten()# 這是neighbor_embeddings shape = [600*10, 172]neighbor_embeddings = self.compute_embedding(memory,neighbors, # 這里有6000個np.repeat(timestamps, n_neighbors), # 也是6000n_layers=n_layers - 1,n_neighbors=n_neighbors)effective_n_neighbors = n_neighbors if n_neighbors > 0 else 1# 這是neighbor_embeddings shape = [600, 10, 172]neighbor_embeddings = neighbor_embeddings.view(len(source_nodes), effective_n_neighbors, -1)# edge_time_embeddings shape is [600, 10, 172]edge_time_embeddings = self.time_encoder(edge_deltas_torch)# self.edge_features shape [157475, 172]# edge_idxs shape [600, 10]# edge_features shape [600, 10, 172]edge_features = self.edge_features[edge_idxs, :]mask = neighbors_torch == 0# 這個聚合在下面"""n_layers: 1source_node_conv_embeddings: 一開始那600個節點的編碼source_nodes_time_embedding: 數據是和timestamps_torch一樣的0矩陣[3*200, 1, 172]neighbor_embeddings: 之前那600個節點的發生過操作的鄰居edge_time_embeddings: 時間差編碼edge_features: 一開始那600個節點，對應的十個鄰居，分別邊的特征是多少mask = [600*10]"""source_embedding = self.aggregate(n_layers, source_node_conv_embeddings,source_nodes_time_embedding,neighbor_embeddings,edge_time_embeddings,edge_features,mask)return source_embeddingdef aggregate(self, n_layers, source_node_features, source_nodes_time_embedding,neighbor_embeddings,edge_time_embeddings, edge_features, mask):return NotImplementedclass GraphAttentionEmbedding(GraphEmbedding):def __init__(self, node_features, edge_features, memory, neighbor_finder, time_encoder, n_layers,n_node_features, n_edge_features, n_time_features, embedding_dimension, device,n_heads=2, dropout=0.1, use_memory=True):super(GraphAttentionEmbedding, self).__init__(node_features, edge_features, memory,neighbor_finder, time_encoder, n_layers,n_node_features, n_edge_features,n_time_features,embedding_dimension, device,n_heads, dropout,use_memory)self.attention_models = torch.nn.ModuleList([TemporalAttentionLayer(n_node_features=n_node_features,n_neighbors_features=n_node_features,n_edge_features=n_edge_features,time_dim=n_time_features,n_head=n_heads,dropout=dropout,output_dimension=n_node_features)for _ in range(n_layers)])def aggregate(self, n_layer, source_node_features, source_nodes_time_embedding,neighbor_embeddings,edge_time_embeddings, edge_features, mask):attention_model = self.attention_models[n_layer - 1]source_embedding, _ = attention_model(source_node_features,source_nodes_time_embedding,neighbor_embeddings,edge_time_embeddings,edge_features,mask)return source_embedding
TimeEncode
class TimeEncode(torch.nn.Module):# Time Encoding proposed by TGATdef __init__(self, dimension):super(TimeEncode, self).__init__()self.dimension = dimension # 172self.w = torch.nn.Linear(1, dimension)# todoself.w.weight = torch.nn.Parameter((torch.from_numpy(1 / 10 ** np.linspace(0, 9, dimension))).float().reshape(dimension, -1))self.w.bias = torch.nn.Parameter(torch.zeros(dimension).float())def forward(self, t): # -> [batch_size, seq_len, dimension]# t has shape [batch_size, seq_len]# Add dimension at the end to apply linear layer --> [batch_size, seq_len, 1]t = t.unsqueeze(dim=2)# output has shape [batch_size, seq_len, dimension]output = torch.cos(self.w(t))return output
NeighborFinder
class NeighborFinder:def __init__(self, adj_list, uniform=False, seed=None):self.node_to_neighbors = []self.node_to_edge_idxs = []self.node_to_edge_timestamps = []for neighbors in adj_list:# Neighbors is a list of tuples (neighbor, edge_idx, timestamp)# We sort the list based on timestampsorted_neighhbors = sorted(neighbors, key=lambda x: x[2])self.node_to_neighbors.append(np.array([x[0] for x in sorted_neighhbors])) # 是一個二維數組，第一個維度表示的是某一個節點，第二個維度表示的是這個節點和那些節點發生的聯系self.node_to_edge_idxs.append(np.array([x[1] for x in sorted_neighhbors]))self.node_to_edge_timestamps.append(np.array([x[2] for x in sorted_neighhbors]))self.uniform = uniformif seed is not None:self.seed = seedself.random_state = np.random.RandomState(self.seed)def find_before(self, src_idx, cut_time):"""Extracts all the interactions happening before cut_time for user src_idx in the overall interaction graph. The returned interactions are sorted by time.Returns 3 lists: neighbors, edge_idxs, timestamps"""i = np.searchsorted(self.node_to_edge_timestamps[src_idx], cut_time)return self.node_to_neighbors[src_idx][:i], self.node_to_edge_idxs[src_idx][:i], self.node_to_edge_timestamps[src_idx][:i]def get_temporal_neighbor(self, source_nodes, timestamps, n_neighbors=20):"""Given a list of users ids and relative cut times, extracts a sampled temporal neighborhood of each user in the list.Params------src_idx_l: List[int]cut_time_l: List[float],num_neighbors: int"""assert (len(source_nodes) == len(timestamps))tmp_n_neighbors = n_neighbors if n_neighbors > 0 else 1# NB! All interactions described in these matrices are sorted in each row by timeneighbors = np.zeros((len(source_nodes), tmp_n_neighbors)).astype( # shape [600, 10]np.int32) # each entry in position (i,j) represent the id of the item targeted by user src_idx_l[i] with an interaction happening before cut_time_l[i]edge_times = np.zeros((len(source_nodes), tmp_n_neighbors)).astype(np.float32) # each entry in position (i,j) represent the timestamp of an interaction between user src_idx_l[i] and item neighbors[i,j] happening before cut_time_l[i]edge_idxs = np.zeros((len(source_nodes), tmp_n_neighbors)).astype(np.int32) # each entry in position (i,j) represent the interaction index of an interaction between user src_idx_l[i] and item neighbors[i,j] happening before cut_time_l[i]for i, (source_node, timestamp) in enumerate(zip(source_nodes, timestamps)):source_neighbors, source_edge_idxs, source_edge_times = self.find_before(source_node,timestamp) # extracts all neighbors, interactions indexes and timestamps of all interactions of user source_node happening before cut_timeif len(source_neighbors) > 0 and n_neighbors > 0:if self.uniform: # if we are applying uniform sampling, shuffles the data above before samplingsampled_idx = np.random.randint(0, len(source_neighbors), n_neighbors)neighbors[i, :] = source_neighbors[sampled_idx]edge_times[i, :] = source_edge_times[sampled_idx]edge_idxs[i, :] = source_edge_idxs[sampled_idx]# re-sort based on timepos = edge_times[i, :].argsort()neighbors[i, :] = neighbors[i, :][pos]edge_times[i, :] = edge_times[i, :][pos]edge_idxs[i, :] = edge_idxs[i, :][pos]else:# Take most recent interactionssource_edge_times = source_edge_times[-n_neighbors:]source_neighbors = source_neighbors[-n_neighbors:]source_edge_idxs = source_edge_idxs[-n_neighbors:]assert (len(source_neighbors) <= n_neighbors)assert (len(source_edge_times) <= n_neighbors)assert (len(source_edge_idxs) <= n_neighbors)neighbors[i, n_neighbors - len(source_neighbors):] = source_neighborsedge_times[i, n_neighbors - len(source_edge_times):] = source_edge_timesedge_idxs[i, n_neighbors - len(source_edge_idxs):] = source_edge_idxsreturn neighbors, edge_idxs, edge_times
class TemporalAttentionLayer(torch.nn.Module):"""Temporal attention layer. Return the temporal embedding of a node given the node itself,its neighbors and the edge timestamps."""def __init__(self, n_node_features, n_neighbors_features, n_edge_features, time_dim,output_dimension, n_head=2,dropout=0.1):super(TemporalAttentionLayer, self).__init__()self.n_head = n_headself.feat_dim = n_node_featuresself.time_dim = time_dimself.query_dim = n_node_features + time_dimself.key_dim = n_neighbors_features + time_dim + n_edge_featuresself.merger = MergeLayer(self.query_dim, n_node_features, n_node_features, output_dimension)self.multi_head_target = nn.MultiheadAttention(embed_dim=self.query_dim,kdim=self.key_dim,vdim=self.key_dim,num_heads=n_head,dropout=dropout)def forward(self, src_node_features, src_time_features, neighbors_features,neighbors_time_features, edge_features, neighbors_padding_mask):""""Temporal attention model:param src_node_features: float Tensor of shape [batch_size, n_node_features]:param src_time_features: float Tensor of shape [batch_size, 1, time_dim]:param neighbors_features: float Tensor of shape [batch_size, n_neighbors, n_node_features]:param neighbors_time_features: float Tensor of shape [batch_size, n_neighbors,time_dim]:param edge_features: float Tensor of shape [batch_size, n_neighbors, n_edge_features]:param neighbors_padding_mask: float Tensor of shape [batch_size, n_neighbors]:return:attn_output: float Tensor of shape [1, batch_size, n_node_features]attn_output_weights: [batch_size, 1, n_neighbors]"""# src_node_features_unrolled shape is [600, 1, 172]src_node_features_unrolled = torch.unsqueeze(src_node_features, dim=1)# 將節點特征和時間特征結合在一起# query shape is [600, 1, 172*2]query = torch.cat([src_node_features_unrolled, src_time_features], dim=2)# 鄰居的特征、邊的特征和時間差特征組合在一起 key shape = [600, 10, 516]key = torch.cat([neighbors_features, edge_features, neighbors_time_features], dim=2)# query shape is [1, 600, 344]query = query.permute([1, 0, 2]) # [1, batch_size, num_of_features]# key shape is [10, 600, 516]key = key.permute([1, 0, 2]) # [n_neighbors, batch_size, num_of_features]# 在dim=1的維度下，要是全為True，那么就代表這一行是沒有用的，反之為Falseinvalid_neighborhood_mask = neighbors_padding_mask.all(dim=1, keepdim=True)# neighbors_padding_mask[invalid_neighborhood_mask.squeeze(), 0] = False# print(query.shape, key.shape)attn_output, attn_output_weights = self.multi_head_target(query=query, key=key, value=key,key_padding_mask=neighbors_padding_mask)# mask = torch.unsqueeze(neighbors_padding_mask, dim=2) # mask [B, N, 1]# mask = mask.permute([0, 2, 1])# attn_output, attn_output_weights = self.multi_head_target(q=query, k=key, v=key,# mask=mask)# attn_output shape = [600, 344]# attn_output_weights = [600, 10]attn_output = attn_output.squeeze()attn_output_weights = attn_output_weights.squeeze()# Source nodes with no neighbors have an all zero attention output. The attention output is# then added or concatenated to the original source node features and then fed into an MLP.# This means that an all zero vector is not used.attn_output = attn_output.masked_fill(invalid_neighborhood_mask, 0)attn_output_weights = attn_output_weights.masked_fill(invalid_neighborhood_mask, 0)# Skip connection with temporal attention over neighborhood and the features of the node itself# attn_output = [600, 172]attn_output = self.merger(attn_output, src_node_features)return attn_output, attn_output_weights
MergeLayer
class MergeLayer(torch.nn.Module):def __init__(self, dim1, dim2, dim3, dim4):super().__init__()self.fc1 = torch.nn.Linear(dim1 + dim2, dim3)self.fc2 = torch.nn.Linear(dim3, dim4)self.act = torch.nn.ReLU()torch.nn.init.xavier_normal_(self.fc1.weight)torch.nn.init.xavier_normal_(self.fc2.weight)def forward(self, x1, x2):x = torch.cat([x1, x2], dim=1)h = self.act(self.fc1(x))return self.fc2(h)
![主流MQ [Kafka、RabbitMQ、ZeroMQ、RocketMQ 和 ActiveMQ]](http://pic.xiahunao.cn/主流MQ [Kafka、RabbitMQ、ZeroMQ、RocketMQ 和 ActiveMQ])
)
