pytorch gnn训练显存暴增主因是图数据加载与消息传递导致中间张量非线性膨胀：全图dataloader合并引发边数爆炸、propagate中alpha/x_j/msg等梯度张量累积、pin_memory异常致reserved内存阶梯上升，须用neighborsampler/clusterdata、.detach()、checkpoint及严格no_grad管控。
torch.cuda.memory_allocated() 突然飙升、训练几轮就 CUDA out of memory，不是模型太大，大概率是图结构本身在“吃显存”——邻接矩阵稀疏但不可忽视，节点特征和消息传递过程中的中间张量会随图规模非线性膨胀。
图数据加载时就埋下显存隐患
很多用户用 DataLoader 直接套图数据（如 PyG 的 Data 对象），却没意识到：单个 Data 实例里 edge_index 是长整型、x 是浮点型，但一旦放进 DataLoader 且 batch_size > 1，PyTorch Geometric 会自动调用 collate 合并为一个大图。这个过程不光拼接 edge_index，还会加偏移、补 batch 向量——所有这些中间结果都保留在 GPU 上，且默认不释放。
- 避免方式：对超大图，优先用
ClusterData或NeighborSampler，别用全图DataLoader - 检查点：打印
data.edge_index.shape[1]和len(data.x)，确认 batch 后边数是否爆炸式增长（比如从 10 万跳到 500 万） - 危险操作：把
data.to(device)放在DataLoader外面再传入，会导致整个原始图常驻显存，后续 batch 全部叠加其上
MessagePassing 中的中间激活值越积越多
图卷积本质是 gather → transform → scatter，而 PyG 默认在 propagate() 内部保留所有中间变量用于反向传播。如果模型堆了 5 层 GAT 或 GCN，每层都要缓存 alpha（注意力权重）、x_j（邻居特征副本）、msg（消息张量），这些张量尺寸与边数正相关，很容易比参数本身还占显存。
- 关键动作：在
forward中对不需要梯度的中间量显式调用.detach()，例如alpha = alpha.detach()（若只用于 masking） - 更稳妥做法：用
torch.utils.checkpoint.checkpoint()包裹conv_layer(x, edge_index)，牺牲一点速度换显存 - 注意陷阱：不要在
propagate()外部手动删msg，PyG 的 autograd 机制可能已绑定引用；应改用checkpoint或重写message()返回轻量结构
图采样器没关 pin_memory 反而加剧显存泄漏
NeighborSampler 或 ClusterData 本意是降显存，但如果设置 pin_memory=True（尤其在多 worker 场景下），CPU 端 pinned memory 会暴涨，间接导致 CUDA 上下文无法及时回收显存——表现为 torch.cuda.memory_reserved() 持续上涨，但 allocated 看似正常。
- 必须关闭：
DataLoader(..., pin_memory=False, num_workers=0)是最稳组合；若需多 worker，务必设pin_memory=False - 验证方法：在训练 loop 开头加
print(torch.cuda.memory_reserved()/1024**2)，看是否阶梯式上升 - 隐藏风险：某些图库（如 DGL）的
GraphDataLoader默认启用pin_memory，需显式覆盖
图模型里的 torch.no_grad() 容易漏掉验证阶段
图任务常需在验证集上跑 full-batch 推理（比如 node classification 的 test accuracy），此时若忘记 with torch.no_grad():，propagate() 仍会构建完整计算图，而验证图往往和训练图一样大——等于白占一份显存，且不释放。
- 典型错误：写
model.eval()却没配no_grad，model.eval()只影响 dropout/bn，不影响 autograd - 安全写法：验证 loop 必须包在
with torch.no_grad():里，且所有输出（如 logits）立刻转.cpu().numpy()或.item() - 特别注意：用
torchmetrics累计指标时，别直接传preds张量，要用preds.detach().cpu()，否则update()内部会隐式持有图
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
