GAT图神经网络代码

以下是一个简单的GAT（Graph Attention Network）图神经网络的代码示例：

import torch
import torch.nn as nn
import torch.nn.functional as F

class GATLayer(nn.Module):
    def __init__(self, in_features, out_features, dropout=0.6, alpha=0.2):
        super(GATLayer, self).__init__()
        self.dropout = dropout
        self.alpha = alpha
        
        self.W = nn.Linear(in_features, out_features, bias=False)
        self.a = nn.Linear(2*out_features, 1, bias=False)
        
    def forward(self, X, adj_matrix):
        h = self.W(X)
        N = h.size(0)
        
        a_input = torch.cat([h.repeat(1, N).view(N*N, -1), h.repeat(N, 1)], dim=1).view(N, -1, 2*h.size(1))
        e = F.leaky_relu(self.a(a_input).squeeze(2), negative_slope=self.alpha)
        
        zero_vec = -9e15*torch.ones_like(e)
        attention = torch.where(adj_matrix > 0, e, zero_vec)
        
        attention = F.softmax(attention, dim=1)
        attention = F.dropout(attention, p=self.dropout, training=self.training)
        
        h_prime = torch.matmul(attention, h)
        
        return F.elu(h_prime)

class GAT(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, num_layers, dropout=0.6, alpha=0.2):
        super(GAT, self).__init__()
        self.hidden_features = hidden_features
        self.num_layers = num_layers
        
        self.layers = nn.ModuleList([GATLayer(in_features, hidden_features, dropout=dropout, alpha=alpha)])
        self.layers.extend([GATLayer(hidden_features, hidden_features, dropout=dropout, alpha=alpha) for _ in range(num_layers-2)])
        self.layers.append(GATLayer(hidden_features, out_features, dropout=dropout, alpha=alpha))
        
    def forward(self, X, adj_matrix):
        h = X
        for layer in self.layers:
            h = layer(h, adj_matrix)
        
        return h

这是一个简单的GAT图神经网络的实现，包括了GATLayer和GAT两个类。GATLayer定义了一个GAT层的操作，GAT则将多个GAT层串联起来构成整个图神经网络。其中，in_features表示输入特征的维度，hidden_features表示隐层特征的维度，out_features表示输出特征的维度，num_layers表示GAT层数，dropout表示dropout率，alpha表示LeakyReLU的斜率。

希望这个代码示例对你有帮助！如有任何问题，请随时提问。