(Hello world级任务)

数据集:Cora

图神经网络领域的标准数据集。

  • 节点 (Nodes):2708 篇学术论文。

  • 边 (Edges):5429 条引用关系(A 引用 B,则 A 和 B 相连)。

  • 特征 (Features):每篇论文的词袋模型向量 (1433维)。

  • 标签 (Labels):论文所属的 7 个学术领域(如 AI, ML 等)。

只看特征,用 MLP 也能分类。但 GNN 的核心优势在于利用关系:

假设属于同一个领域的论文,倾向于互相引用。GNN 通过聚合邻居的信息,能让分类更准确。这在图不够稠密或特征有噪音时尤为重要。

搭建一个包含两层卷积的 GCN 模型。GCN 的核心机制是消息传递:每个节点收集周围邻居的特征,取平均后更新自己。

使用kaggle notebook运行,先安装包

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import os
import torch

# 自动检测当前的 PyTorch 和 CUDA 版本
torch_ver = torch.__version__.split('+')[0]  # 例如 2.1.2
cuda_ver = torch.version.cuda.replace('.', '') # 例如 118
print(f"检测到环境: PyTorch {torch_ver}, CUDA {cuda_ver}")

# 构造正确的下载链接
url = f"https://data.pyg.org/whl/torch-{torch_ver}+cu{cuda_ver}.html"

# 安装 
print(f"正在从 {url} 下载安装...")
!pip install -q torch_geometric
!pip install -q torch_scatter torch_sparse -f {url}

引入数据

import torch
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T

# 1. 加载 Cora 数据集
# PyG 会自动帮你下载并处理成图格式
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]  # Cora 只有一张大图,所以取第0个元素

# 2. 统计学视角的维度检查
print(f'数据集: {dataset}:')
print('======================')
print(f'节点数量 (Number of papers): {data.num_nodes}')
print(f'特征维度 (Number of words): {data.num_features}')
print(f'类别数量 (Number of classes): {dataset.num_classes}')
print(f'边数量 (Number of edges): {data.num_edges}')
print(f'平均节点度 (Average node degree): {data.num_edges / data.num_nodes:.2f}')

print('\n数据对象概览:')
print(data)

搭建GCN

import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 定义两层 GCN
        # 第一层:输入特征维度 -> 隐藏层 (比如16)
        self.conv1 = GCNConv(dataset.num_features, 32)
        # 第二层:隐藏层 -> 输出类别 (7类)
        self.conv2 = GCNConv(32, dataset.num_classes)

    def forward(self, data):
        # 取出特征矩阵 x 和 邻接关系 edge_index
        x, edge_index = data.x, data.edge_index

        # --- 第一层卷积 ---
        # 这里的计算公式隐含了 A_hat * X * W (邻接矩阵归一化乘特征乘权重)
        x = self.conv1(x, edge_index)
        x = F.relu(x) # 激活函数
        x = F.dropout(x, training=self.training) # 防止过拟合

        # --- 第二层卷积 ---
        x = self.conv2(x, edge_index)
        
        # 输出层用 log_softmax 做多分类
        return F.log_softmax(x, dim=1)

# 初始化模型并移到 GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN().to(device)
data = data.to(device) # 把数据也搬到 GPU 上
print(model)

训练

optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# 训练函数
def train():
    model.train()
    optimizer.zero_grad()
    
    # 1. 前向传播
    out = model(data) 
    
    # 2. 计算损失
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    
    # 3. 反向传播
    loss.backward()
    optimizer.step()
    return loss.item()

# 测试函数
@torch.no_grad()  # <--- 这里改好了!没有冒号
def test():
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1) 
    
    # 计算准确率
    correct = pred[data.test_mask] == data.y[data.test_mask]
    acc = int(correct.sum()) / int(data.test_mask.sum())
    return acc

# === 开始跑循环 ===
print("开始训练...")
for epoch in range(1, 201):
    loss = train()
    if epoch % 20 == 0:
        acc = test()
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Test Acc: {acc:.4f}')

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐