从零到一：用PyTorch Geometric实现你的第一个GraphSAGE模型（附完整代码）
从零到一：用PyTorch Geometric实现你的第一个GraphSAGE模型（附完整代码）
第一次接触图神经网络时，我被它的独特魅力所吸引——它能够直接处理社交网络、分子结构这类非欧几里得数据。但真正动手实现时，却遇到了各种工程难题：如何高效处理图数据？邻居采样该怎么实现？模型训练为什么总是不收敛？本文将带你从零开始，用PyTorch Geometric（PyG）这个利器，一步步构建可运行的GraphSAGE模型。
1. 环境准备与数据加载
1.1 安装PyTorch Geometric
PyG是图神经网络领域的瑞士军刀，但它的安装有些特殊技巧。推荐使用conda创建虚拟环境：
conda create -n graphsage python=3.8 conda activate graphsage pip install torch torchvision torchaudio pip install torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-1.10.0+cu113.html pip install torch-geometric注意：torch-geometric需要与PyTorch版本严格匹配，建议先查看官方安装指南。如果遇到C++扩展编译错误，可以尝试安装预编译版本。
1.2 加载Cora数据集
让我们用经典的Cora论文引用网络作为示例：
from torch_geometric.datasets import Planetoid dataset = Planetoid(root='/tmp/Cora', 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}')这个数据集包含2708篇论文节点，每篇论文有1433维的词袋特征，边代表引用关系。我们可以用以下代码可视化节点特征分布：
import matplotlib.pyplot as plt from sklearn.manifold import TSNE def visualize(h, color): z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy()) plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap="Set2") plt.show() visualize(data.x, data.y)2. GraphSAGE模型架构解析
2.1 邻居聚合机制
GraphSAGE的核心在于它的多层聚合机制。与GCN不同，它支持多种聚合方式：
| 聚合类型 | 公式 | 特点 |
|---|---|---|
| Mean | $\frac{1}{ | N(v) |
| LSTM | LSTM([h_u, ∀u∈N(v)]) | 考虑邻居顺序，需随机排列 |
| Pool | max(σ(W_poolh_u+b)) | 非线性变换后取最大 |
2.2 PyG实现方案
PyG提供了SAGEConv层，我们只需关注网络设计：
import torch import torch.nn.functional as F from torch_geometric.nn import SAGEConv class GraphSAGE(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels): super().__init__() self.conv1 = SAGEConv(in_channels, hidden_channels, aggr='mean') self.conv2 = SAGEConv(hidden_channels, out_channels, aggr='mean') def forward(self, x, edge_index): x = self.conv1(x, edge_index) x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)这个两层的网络已经能处理大多数任务。如果想尝试不同聚合方式，只需修改aggr参数：
self.conv1 = SAGEConv(in_channels, hidden_channels, aggr='lstm')3. 训练与评估实战
3.1 训练流程优化
标准的训练循环需要特别注意图数据的特殊性：
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = GraphSAGE(dataset.num_features, 16, dataset.num_classes).to(device) data = data.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) def train(): model.train() optimizer.zero_grad() out = model(data.x, data.edge_index) loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step() return loss.item()提示：图数据通常存在类别不平衡问题，可以尝试在损失函数中加入类别权重：
class_weight = 1. / torch.bincount(data.y[data.train_mask]) criterion = torch.nn.NLLLoss(weight=class_weight)
3.2 邻居采样技巧
全图训练在大规模图上不现实。PyG的NeighborSampler可以实现高效采样：
from torch_geometric.loader import NeighborSampler train_loader = NeighborSampler(data.edge_index, node_idx=data.train_mask, sizes=[10, 5], batch_size=256, shuffle=True) def sampled_train(): model.train() total_loss = 0 for batch_size, n_id, adjs in train_loader: adjs = [adj.to(device) for adj in adjs] optimizer.zero_grad() out = model(data.x[n_id].to(device), adjs) loss = F.nll_loss(out, data.y[n_id[:batch_size]].to(device)) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(train_loader)采样参数sizes=[10,5]表示：第一层采样10个邻居，第二层从这10个节点各采样5个邻居。
4. 高级技巧与性能调优
4.1 特征工程增强
原始节点特征可能不够丰富，可以尝试：
- 特征标准化：
from sklearn.preprocessing import StandardScaler scaler = StandardScaler() data.x = torch.tensor(scaler.fit_transform(data.x.numpy()), dtype=torch.float)- 添加结构特征：
degree = torch_geometric.utils.degree(data.edge_index[0]) data.x = torch.cat([data.x, degree.view(-1, 1)], dim=1)4.2 模型深度与过拟合
增加网络深度时要注意：
- 使用残差连接防止梯度消失：
class ResGraphSAGE(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, num_layers=3): super().__init__() self.convs = torch.nn.ModuleList() self.convs.append(SAGEConv(in_channels, hidden_channels)) for _ in range(num_layers - 2): self.convs.append(SAGEConv(hidden_channels, hidden_channels)) self.convs.append(SAGEConv(hidden_channels, out_channels)) def forward(self, x, edge_index): for i, conv in enumerate(self.convs[:-1]): x = conv(x, edge_index) x = F.relu(x) x = F.dropout(x, p=0.5, training=self.training) return self.convs[-1](x, edge_index)- 早停法监控验证集性能：
best_val_acc = 0 patience = 20 counter = 0 for epoch in range(1, 201): loss = train() val_acc = test(data.val_mask) if val_acc > best_val_acc: best_val_acc = val_acc counter = 0 else: counter += 1 if counter == patience: print(f'Early stopping at epoch {epoch}') break4.3 可视化分析
理解模型行为的关键是观察节点嵌入的变化：
def visualize_progress(model, data, epoch): model.eval() with torch.no_grad(): out = model(data.x, data.edge_index) visualize(out, data.y) plt.title(f'Epoch {epoch}') plt.savefig(f'embedding_{epoch}.png') for epoch in range(1, 51): train() if epoch % 10 == 0: visualize_progress(model, data, epoch)这个可视化能清晰展示模型如何逐步将同类节点聚集在一起。