图卷积网络（GCN）简单示例

# 导入必要的库 import torch import torch.nn.functional as F from torch_geometric.nn import GCNConv from torch_geometric.datasets import Planetoid import matplotlib.pyplot as plt import networkx as nx from torch_geometric.utils import to_networkx # 1. 加载数据集（使用Cora数据集，这是一个引用网络数据集） dataset = Planetoid(root='/tmp/Cora', name='Cora') data = dataset[0] # 获取图数据 # 2. 定义GCN模型 class GCN(torch.nn.Module): def __init__(self): super(GCN, self).__init__() # 定义两层GCN卷积层 self.conv1 = GCNConv(dataset.num_node_features, 16) self.conv2 = GCNConv(16, dataset.num_classes) def forward(self, data): x, edge_index = data.x, data.edge_index # 第一层卷积+ReLU激活函数 x = self.conv1(x, edge_index) x = F.relu(x) # 第二层卷积+Softmax输出 x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1) # 3. 初始化模型和优化器 model = GCN() optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) # 4. 定义绘制节点图函数 def plot_graph(data, color_map=None, title="Graph"): # 将数据转换为NetworkX图 G = to_networkx(data, to_undirected=True) plt.figure(figsize=(8, 8)) # 绘制图，并为节点上色 nx.draw(G, pos=nx.spring_layout(G), with_labels=False, node_color=color_map, node_size=50, cmap="coolwarm") plt.title(title) plt.show() # 使用真实标签颜色绘制原始图 color_map = data.y.numpy() plot_graph(data, color_map, title="Original Graph with True Labels") # 5. 训练模型 model.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() if epoch % 20 == 0: print(f'Epoch {epoch}, Loss: {loss.item()}') # 6. 评估模型，并可视化预测结果 model.eval() _, pred = model(data).max(dim=1) # 使用预测标签颜色绘制图 pred_color_map = pred.numpy() # 使用预测标签作为颜色映射 plot_graph(data, pred_color_map, title="Graph with Predicted Labels") # 计算并输出准确率 correct = (pred[data.test_mask] == data.y[data.test_mask]).sum() accuracy = int(correct) / int(data.test_mask.sum()) print(f'Accuracy: {accuracy:.4f}')