图注意网力络论文详解和PyTorch实现

 import torch
 from torch import nn
 
 # in_features -> F and out_feature -> F'
 in_features = ...
 out_feature = ...
 
 # instanciate the learnable weight matrix W (FxF')
 W = nn.Parameter(torch.empty(size=(in_features, out_feature)))
 
 #  Initialize the weight matrix W
 nn.init.xavier_normal_(W)
 
 # multiply W and h (h is input features of all the nodes -> NxF matrix)
 h_transformed = torch.mm(h, W)

 # instanciate the learnable attention parameter vector `a`
 a = nn.Parameter(torch.empty(size=(2 * out_feature, 1)))
 
 # Initialize the parameter vector `a`
 nn.init.xavier_normal_(a)
 
 # we obtained `h_transformed` in the previous code snippet
 
 # calculating the dot product of all node embeddings
 # and first half the attention vector parameters (corresponding to neighbor messages)
 source_scores = torch.matmul(h_transformed, self.a[:out_feature, :])
 
 # calculating the dot product of all node embeddings
 # and second half the attention vector parameters (corresponding to target node)
 target_scores = torch.matmul(h_transformed, self.a[out_feature:, :])
 
 # broadcast add 
 e = source_scores + target_scores.T
 e = self.leakyrelu(e)

 connectivity_mask = -9e16 * torch.ones_like(e)
 # adj_mat is the N by N adjacency matrix
 e = torch.where(adj_mat > 0, e, connectivity_mask) # masked attention scores
         
 # attention coefficients are computed as a softmax over the rows
 # for each column j in the attention score matrix e
 attention = F.softmax(e, dim=-1)

 # final node embeddings are computed as a weighted average of the features of its neighbors
 h_prime = torch.matmul(attention, h_transformed)

 import torch
 from torch import nn
 import torch.nn.functional as F
 
 ################################
 ###  GAT LAYER DEFINITION    ###
 ################################
 
 class GraphAttentionLayer(nn.Module):
 
     def __init__(self, in_features: int, out_features: int,
                  n_heads: int, concat: bool = False, dropout: float = 0.4,
                  leaky_relu_slope: float = 0.2):
         super(GraphAttentionLayer, self).__init__()
 
         self.n_heads = n_heads # Number of attention heads
         self.concat = concat # wether to concatenate the final attention heads
         self.dropout = dropout # Dropout rate
 
         if concat: # concatenating the attention heads
             self.out_features = out_features # Number of output features per node
             assert out_features % n_heads == 0 # Ensure that out_features is a multiple of n_heads
             self.n_hidden = out_features // n_heads
         else: # averaging output over the attention heads (Used in the main paper)
             self.n_hidden = out_features
 
         #  A shared linear transformation, parametrized by a weight matrix W is applied to every node
         #  Initialize the weight matrix W 
         self.W = nn.Parameter(torch.empty(size=(in_features, self.n_hidden * n_heads)))
 
         # Initialize the attention weights a
         self.a = nn.Parameter(torch.empty(size=(n_heads, 2 * self.n_hidden, 1)))
 
         self.leakyrelu = nn.LeakyReLU(leaky_relu_slope) # LeakyReLU activation function
         self.softmax = nn.Softmax(dim=1) # softmax activation function to the attention coefficients
 
         self.reset_parameters() # Reset the parameters
 
 
     def reset_parameters(self):
 
         nn.init.xavier_normal_(self.W)
         nn.init.xavier_normal_(self.a)
 
     def _get_attention_scores(self, h_transformed: torch.Tensor):
         
         source_scores = torch.matmul(h_transformed, self.a[:, :self.n_hidden, :])
         target_scores = torch.matmul(h_transformed, self.a[:, self.n_hidden:, :])
 
         # broadcast add 
         # (n_heads, n_nodes, 1) + (n_heads, 1, n_nodes) = (n_heads, n_nodes, n_nodes)
         e = source_scores + target_scores.mT
         return self.leakyrelu(e)
 
     def forward(self,  h: torch.Tensor, adj_mat: torch.Tensor):
 
         n_nodes = h.shape[0]
 
         # Apply linear transformation to node feature -> W h
         # output shape (n_nodes, n_hidden * n_heads)
         h_transformed = torch.mm(h, self.W)
         h_transformed = F.dropout(h_transformed, self.dropout, training=self.training)
 
         # splitting the heads by reshaping the tensor and putting heads dim first
         # output shape (n_heads, n_nodes, n_hidden)
         h_transformed = h_transformed.view(n_nodes, self.n_heads, self.n_hidden).permute(1, 0, 2)
         
         # getting the attention scores
         # output shape (n_heads, n_nodes, n_nodes)
         e = self._get_attention_scores(h_transformed)
 
         # Set the attention score for non-existent edges to -9e15 (MASKING NON-EXISTENT EDGES)
         connectivity_mask = -9e16 * torch.ones_like(e)
         e = torch.where(adj_mat > 0, e, connectivity_mask) # masked attention scores
         
         # attention coefficients are computed as a softmax over the rows
         # for each column j in the attention score matrix e
         attention = F.softmax(e, dim=-1)
         attention = F.dropout(attention, self.dropout, training=self.training)
 
         # final node embeddings are computed as a weighted average of the features of its neighbors
         h_prime = torch.matmul(attention, h_transformed)
 
         # concatenating/averaging the attention heads
         # output shape (n_nodes, out_features)
         if self.concat:
             h_prime = h_prime.permute(1, 0, 2).contiguous().view(n_nodes, self.out_features)
         else:
             h_prime = h_prime.mean(dim=0)
 
         return h_prime

 class GAT(nn.Module):
 
     def __init__(self,
         in_features,
         n_hidden,
         n_heads,
         num_classes,
         concat=False,
         dropout=0.4,
         leaky_relu_slope=0.2):
 
         super(GAT, self).__init__()
 
         # Define the Graph Attention layers
         self.gat1 = GraphAttentionLayer(
             in_features=in_features, out_features=n_hidden, n_heads=n_heads,
             concat=concat, dropout=dropout, leaky_relu_slope=leaky_relu_slope
             )
         
         self.gat2 = GraphAttentionLayer(
             in_features=n_hidden, out_features=num_classes, n_heads=1,
             concat=False, dropout=dropout, leaky_relu_slope=leaky_relu_slope
             )
 
     def forward(self, input_tensor: torch.Tensor , adj_mat: torch.Tensor):
 
 
         # Apply the first Graph Attention layer
         x = self.gat1(input_tensor, adj_mat)
         x = F.elu(x) # Apply ELU activation function to the output of the first layer
 
         # Apply the second Graph Attention layer
         x = self.gat2(x, adj_mat)
 
         return F.softmax(x, dim=1) # Apply softmax activation function