03-图神经网络（GNN）¶

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

class GCNLayer(nn.Module):  # 继承nn.Module定义神经网络层
    """单层图卷积"""
    def __init__(self, in_features, out_features):  # __init__构造方法，创建对象时自动调用
        super().__init__()  # super()调用父类方法
        self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
        self.bias = nn.Parameter(torch.FloatTensor(out_features))
        nn.init.xavier_uniform_(self.weight)
        nn.init.zeros_(self.bias)

    def forward(self, x, adj_norm):
        """
        x: 节点特征矩阵 [N, in_features]
        adj_norm: 归一化邻接矩阵 D^{-1/2} A~ D^{-1/2}
        """
        support = torch.mm(x, self.weight)       # 线性变换
        output = torch.spmm(adj_norm, support)    # 邻居聚合
        return output + self.bias

class GCN(nn.Module):
    def __init__(self, n_feat, n_hidden, n_class, dropout=0.5):
        super().__init__()
        self.gc1 = GCNLayer(n_feat, n_hidden)
        self.gc2 = GCNLayer(n_hidden, n_class)
        self.dropout = dropout

    def forward(self, x, adj):
        x = F.relu(self.gc1(x, adj))
        x = F.dropout(x, self.dropout, training=self.training)
        x = self.gc2(x, adj)
        return F.log_softmax(x, dim=1)

class GATLayer(nn.Module):
    def __init__(self, in_features, out_features, n_heads=4, concat=True):
        super().__init__()
        self.n_heads = n_heads
        self.concat = concat
        self.W = nn.Linear(in_features, out_features * n_heads, bias=False)
        self.a = nn.Parameter(torch.FloatTensor(n_heads, 2 * out_features))
        self.leaky_relu = nn.LeakyReLU(0.2)
        nn.init.xavier_uniform_(self.a)
        self.out_features = out_features

    def forward(self, x, edge_index):
        """
        x: [N, in_features]
        edge_index: [2, E] 边索引
        """
        N = x.size(0)
        h = self.W(x).view(N, self.n_heads, self.out_features)  # [N, heads, out]  # 链式调用，连续执行多个方法

        src, dst = edge_index  # 源节点与目标节点
        h_src = h[src]  # [E, heads, out]
        h_dst = h[dst]  # [E, heads, out]

        # 计算注意力系数
        cat_feat = torch.cat([h_src, h_dst], dim=-1)  # [E, heads, 2*out]
        e = (cat_feat * self.a.unsqueeze(0)).sum(dim=-1)  # [E, heads]  # unsqueeze增加一个维度
        e = self.leaky_relu(e)

        # softmax归一化（针对每个目标节点的所有入边）
        # 使用scatter实现per-node softmax：对每个目标节点，其所有入边的注意力系数归一化
        e_max = e.new_full((N, self.n_heads), -1e9)
        e_max.scatter_reduce_(0, dst.unsqueeze(-1).expand_as(e), e, reduce='amax')
        e_stable = torch.exp(e - e_max[dst])  # 数值稳定
        e_sum = torch.zeros(N, self.n_heads, device=x.device)
        e_sum.scatter_add_(0, dst.unsqueeze(-1).expand_as(e_stable), e_stable)
        alpha = e_stable / (e_sum[dst] + 1e-16)  # [E, heads]

        # 聚合
        out = torch.zeros(N, self.n_heads, self.out_features, device=x.device)
        out.scatter_add_(0, dst.unsqueeze(-1).unsqueeze(-1).expand_as(h_src * alpha.unsqueeze(-1)),
                         h_src * alpha.unsqueeze(-1))

        if self.concat:
            return out.view(N, -1)  # [N, heads * out]
        else:
            return out.mean(dim=1)  # [N, out]

# DiffPool 核心思想：学习软聚类分配矩阵
class DiffPoolLayer(nn.Module):
    def __init__(self, in_feat, n_clusters):
        super().__init__()
        self.gnn_embed = GCNLayer(in_feat, in_feat)   # 特征GNN
        self.gnn_pool = GCNLayer(in_feat, n_clusters)  # 池化GNN

    def forward(self, x, adj):
        z = F.relu(self.gnn_embed(x, adj))    # 节点嵌入
        s = F.softmax(self.gnn_pool(x, adj), dim=-1)  # 分配矩阵 [N, K]

        # 粗化
        x_pooled = s.T @ z      # [K, feat]
        adj_pooled = s.T @ adj @ s  # [K, K]
        return x_pooled, adj_pooled

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv
from torch_geometric.datasets import Planetoid

# 加载 Cora 数据集
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]

print(f"节点数: {data.num_nodes}")
print(f"边数: {data.num_edges}")
print(f"特征维度: {data.num_node_features}")
print(f"类别数: {dataset.num_classes}")

class GATNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GATConv(dataset.num_node_features, 8, heads=8)
        self.conv2 = GATConv(64, dataset.num_classes, heads=1, concat=False)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

# 训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GATNet().to(device)  # .to(device)将数据移至GPU/CPU
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

model.train()  # train()开启训练模式
for epoch in range(200):
    optimizer.zero_grad()  # 清零梯度，防止梯度累积
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()  # 反向传播计算梯度
    optimizer.step()  # 根据梯度更新模型参数

# 测试
model.eval()
pred = model(data).argmax(dim=1)
correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
acc = int(correct) / int(data.test_mask.sum())
print(f"测试准确率: {acc:.4f}")