GAT代码示例

以下是一个简单的GAT代码示例，用于节点分类任务：

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

class GATLayer(nn.Module):
    def __init__(self, in_dim, out_dim):
        super(GATLayer, self).__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.W = nn.Linear(in_dim, out_dim)
        self.a = nn.Linear(2 * out_dim, 1)

    def forward(self, h, adj):
        Wh = self.W(h)
        a_input = self._prepare_attentional_mechanism_input(Wh)
        e = F.leaky_relu(self.a(a_input))
        attention = F.softmax(e, dim=1)
        h_prime = torch.matmul(attention.transpose(1, 2), Wh).squeeze(1)
        return h_prime

    def _prepare_attentional_mechanism_input(self, Wh):
        N = Wh.size()[0]
        Wh_repeated_in_chunks = Wh.repeat_interleave(N, dim=0)
        Wh_repeated_alternating = Wh.repeat(N, 1)
        all_combinations_matrix = torch.cat([Wh_repeated_in_chunks, Wh_repeated_alternating], dim=1)
        return all_combinations_matrix.view(N, N, 2 * self.out_dim)

class GATNet(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super(GATNet, self).__init__()
        self.layer1 = GATLayer(in_dim, hidden_dim)
        self.layer2 = GATLayer(hidden_dim, out_dim)

    def forward(self, x, adj):
        x = F.dropout(x, 0.5, training=self.training)
        x = F.elu(self.layer1(x, adj))
        x = F.dropout(x, 0.5, training=self.training)
        x = self.layer2(x, adj)
        return F.log_softmax(x, dim=1)

这里定义了一个GAT层和一个GAT网络。GAT层包含一个线性变换和一个注意力机制，用于计算节点之间的重要性。GAT网络由两个GAT层组成，用于将节点嵌入分类空间。在前向传播过程中，节点特征被输入到网络中，并在每个GAT层之间进行dropout操作。最后，输出通过log_softmax进行归一化。