GAT模型的代码

以下是一个简单的GAT模型的代码实现，包括输入层、多头注意力层和输出层：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class GATConv(MessagePassing):
    def __init__(self, in_channels, out_channels, heads=1, concat=True, bias=True, **kwargs):
        super(GATConv, self).__init__(aggr='add', **kwargs)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.concat = concat
        self.weight = nn.Parameter(torch.Tensor(in_channels, heads * out_channels))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(heads * out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.weight)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, x, edge_index):
        x = torch.matmul(x, self.weight)
        return self.propagate(edge_index, x=x)

    def message(self, x_j, edge_index, size_i):
        alpha = (x_j * self.weight.view(self.in_channels, self.heads, self.out_channels)).sum(dim=-1)
        alpha = F.leaky_relu(alpha, negative_slope=0.2)
        alpha = self.softmax(alpha, edge_index[0])
        return x_j * alpha.view(-1, self.heads, 1)

    def update(self, aggr_out):
        if self.concat:
            return aggr_out.view(-1, self.heads * self.out_channels)
        else:
            return aggr_out.mean(dim=1)

class GAT(nn.Module):
    def __init__(self, in_channels, out_channels, hidden_channels, heads=1, depth=1, **kwargs):
        super(GAT, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels
        self.heads = heads
        self.depth = depth
        self.convs = nn.ModuleList()
        self.convs.append(GATConv(in_channels, hidden_channels, heads=heads, **kwargs))
        for i in range(depth - 2):
            self.convs.append(GATConv(hidden_channels * heads, hidden_channels, heads=heads, **kwargs))
        self.convs.append(GATConv(hidden_channels * heads, out_channels, heads=1, concat=False, **kwargs))

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.5, training=self.training)
        for i in range(self.depth - 1):
            x = F.elu(self.convs[i](x, edge_index))
            x = F.dropout(x, p=0.5, training=self.training)
        x = self.convs[-1](x, edge_index)
        return x

其中，`GATConv`是GAT的核心模块，它继承了PyTorch Geometric中的`MessagePassing`类，用于实现图卷积操作。`GAT`则是多层GAT的模型实现。具体实现过程可以参考注释。