PyTorch-Geometric训练模型时输入数据未使用train_mask的相关疑问

核心结论 #

不需要在给模型传入data时主动应用train_mask，这是全图训练GNN的标准写法，官方示例也遵循该逻辑。

具体原因说明 #

• 保证图结构完整性：GNN的核心机制是消息传递，训练集节点的表征计算依赖其邻居节点的特征与邻接关系，哪怕邻居属于验证集或测试集。如果提前通过train_mask过滤掉非训练节点，会直接破坏原始图的结构，导致训练集节点无法拿到完整的邻接信息，最终训练出的模型泛化能力会大幅下降。
• 不会产生额外的反向传播开销：损失计算时仅选择了train_mask对应的节点输出和标签参与计算，反向传播阶段只会对和这部分损失相关的路径计算梯度，非训练节点的前向输出没有参与损失计算，不会产生额外的梯度计算开销。
• 你提到的前向计算浪费确实存在：全图训练时确实会计算所有节点的输出，非训练节点的前向计算属于冗余开销，这是全图训练的固有特性。如果图规模很大，想要避免这部分浪费，可以改用PyG提供的小批次采样训练方案，每次仅采样和当前批次训练节点相关的子图输入模型，就能只计算必要节点的表征。

补充说明 #

你提供的代码写法是完全正确的，不需要修改输入阶段的逻辑，适合中小规模图的训练场景。

内容的提问来源于stack exchange，提问作者Christopher Mills