用DGL和PyTorch复现HAN：手把手教你搞定异构图注意力网络（附完整代码）
用DGL和PyTorch复现HAN：从零实现异构图注意力网络
在现实世界的图数据中，节点和边往往具有多种类型——学术引用网络包含论文、作者、会议等不同实体，电影推荐系统涉及电影、演员、导演等多种对象。这种异构特性使得传统图神经网络难以直接应用。异构图注意力网络(HAN)通过双层注意力机制，不仅解决了异构图的建模难题，还赋予了模型语义理解能力。本文将带您从零开始，用DGL和PyTorch实现这个强大的模型。
1. 环境配置与数据准备
1.1 工具链搭建
确保使用Python 3.8+环境，推荐通过conda创建独立环境：
conda create -n han python=3.8 conda activate han pip install torch==1.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install dgl-cu113==0.9.0注意：CUDA版本需与本地环境匹配，CPU版本可去掉cu113后缀
关键库版本要求：
- PyTorch ≥ 1.10
- DGL ≥ 0.8
- scikit-learn (用于评估指标)
1.2 数据加载与预处理
以IMDB数据集为例，我们需要处理三种节点类型：
import dgl from dgl.data import IMDbDataset # 加载原始数据 dataset = IMDbDataset() graph = dataset[0] # 获取异构图对象 # 节点类型查看 print(graph.ntypes) # ['movie', 'actor', 'director']典型的数据预处理流程包括：
- 特征标准化：对词袋特征做L2归一化
- 元路径定义：确定有意义的连接模式
- 邻居图构建：为每种元路径创建同构子图
import torch.nn.functional as F # 特征归一化 for ntype in graph.ntypes: graph.nodes[ntype].data['feat'] = F.normalize( graph.nodes[ntype].data['feat'], p=2, dim=1) # 定义元路径 metapaths = { 'MAM': ['movie', 'actor', 'movie'], 'MDM': ['movie', 'director', 'movie'] }2. 模型架构深度解析
2.1 节点级注意力实现
节点级注意力是HAN的第一层抽象，其核心是为同一元路径下的邻居分配差异化权重。我们通过NodeLevelAttention模块实现：
import torch import torch.nn as nn import torch.nn.functional as F class NodeLevelAttention(nn.Module): def __init__(self, in_size, out_size): super().__init__() self.project = nn.Sequential( nn.Linear(in_size, out_size), nn.Tanh(), nn.Linear(out_size, 1, bias=False) ) def forward(self, features, neighbors): """ features: 源节点特征 [N, D] neighbors: 邻居特征 [N, K, D] """ # 扩展源节点特征 src_features = features.unsqueeze(1) # [N, 1, D] # 计算注意力分数 attention_scores = self.project( torch.cat([src_features.expand(-1, neighbors.size(1), -1), neighbors], dim=-1) ).squeeze(-1) # [N, K] # 归一化得到注意力权重 return F.softmax(attention_scores, dim=1)关键实现细节：
- 使用两层MLP计算注意力分数
- Tanh激活增强非线性
- 批处理实现高效计算
2.2 语义级注意力设计
语义级注意力是HAN的第二层抽象，用于融合不同元路径的语义信息：
class SemanticAttention(nn.Module): def __init__(self, in_size, hidden_size=128): super().__init__() self.project = nn.Sequential( nn.Linear(in_size, hidden_size), nn.Tanh(), nn.Linear(hidden_size, 1, bias=False) ) def forward(self, embeddings): """ embeddings: 多个元路径的嵌入 [N, P, D] """ # 计算每个元路径的重要性 weights = self.project(embeddings) # [N, P, 1] weights = F.softmax(weights.squeeze(-1), dim=1) # [N, P] # 加权融合 return (embeddings * weights.unsqueeze(-1)).sum(1) # [N, D]3. 完整HAN模型实现
3.1 模型组装
将各组件整合为完整HAN模型：
class HAN(nn.Module): def __init__(self, metapaths, in_size, hidden_size, out_size, num_heads): super().__init__() self.metapaths = metapaths self.num_heads = num_heads # 节点级注意力模块 self.node_attentions = nn.ModuleDict() for mp in metapaths: self.node_attentions[mp] = nn.ModuleList([ NodeLevelAttention(in_size, hidden_size) for _ in range(num_heads) ]) # 语义级注意力模块 self.semantic_attention = SemanticAttention(hidden_size * num_heads) # 输出层 self.predict = nn.Linear(hidden_size * num_heads, out_size) def forward(self, g, h): semantic_embeddings = [] # 对每个元路径处理 for mp, attentions in self.node_attentions.items(): # 获取元路径邻居图 meta_g = dgl.metapath_reachable_graph(g, mp) # 多头注意力 heads = [] for attn in attentions: # 计算注意力权重 weights = attn(h, h[meta_g.edges()]) # 加权聚合 heads.append(torch.matmul(weights, h[meta_g.edges()[1]])) semantic_embeddings.append(torch.cat(heads, dim=1)) # 语义级融合 final_embedding = self.semantic_attention( torch.stack(semantic_embeddings, dim=1) ) return self.predict(final_embedding)3.2 训练流程优化
针对异构图的特性，我们设计专门的训练策略：
def train(model, g, features, labels, train_mask, val_mask, epochs=100): optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=0.001) criterion = nn.CrossEntropyLoss() best_val_acc = 0 for epoch in range(epochs): model.train() logits = model(g, features) loss = criterion(logits[train_mask], labels[train_mask]) optimizer.zero_grad() loss.backward() optimizer.step() # 验证集评估 with torch.no_grad(): model.eval() val_logits = model(g, features) val_pred = val_logits.argmax(1) val_acc = (val_pred[val_mask] == labels[val_mask]).float().mean() if val_acc > best_val_acc: best_val_acc = val_acc torch.save(model.state_dict(), 'best_model.pth') print(f'Epoch {epoch:02d} | Loss: {loss:.4f} | Val Acc: {val_acc:.4f}')4. 实战技巧与问题排查
4.1 常见错误解决方案
在实现过程中，开发者常遇到以下问题：
| 错误类型 | 可能原因 | 解决方案 |
|---|---|---|
| 维度不匹配 | 元路径邻居数量不一致 | 使用masked attention或填充 |
| 梯度消失 | 注意力权重过于集中 | 增加dropout或温度参数 |
| 内存溢出 | 邻居采样过多 | 限制最大邻居数或使用采样 |
4.2 性能优化技巧
邻居采样：对于大规模图，采用随机采样策略
def sample_neighbors(g, nodes, metapath, fanout): edges = g.metapath_random_walk(metapath, nodes, fanout) return torch.unique(edges.flatten())混合精度训练：显著减少显存占用
from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() with autocast(): loss = criterion(model(g, features), labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意力可视化：增强模型可解释性
def visualize_attention(g, model, node_id): # 获取节点级注意力 with torch.no_grad(): model.eval() _, node_attentions = model(g, features, return_attn=True) # 绘制热力图 for mp, attn in node_attentions.items(): plt.figure(figsize=(10,5)) sns.heatmap(attn[node_id].cpu().numpy()) plt.title(f'Attention weights for {mp}') plt.show()
4.3 扩展应用场景
HAN的灵活性使其可应用于多种图学习任务：
- 推荐系统：处理用户-商品-标签异构网络
- 知识图谱：建模实体-关系复杂语义
- 生物网络：分析蛋白质-化合物相互作用
以下是一个推荐系统的改造示例：
class RecommenderHAN(HAN): def __init__(self, metapaths, in_size, hidden_size, num_heads): super().__init__(metapaths, in_size, hidden_size, 1, num_heads) # 修改输出层为评分预测 self.predict = nn.Sequential( nn.Linear(hidden_size * num_heads, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1) ) def forward(self, g, user_nodes, item_nodes): embeddings = super().forward(g) user_emb = embeddings[user_nodes] item_emb = embeddings[item_nodes] return torch.sigmoid((user_emb * item_emb).sum(1))实现过程中发现，对注意力权重施加L2正则能有效防止过拟合，而采用LeakyReLU替代Tanh在节点级注意力中通常能获得更好效果。对于超参数选择，hidden_size设为128、num_heads设为8在大多数场景下表现均衡。