如何在PyTorch Geometric特定层中实现节点掩码与特征聚合？

针对你的需求，PyTorch Geometric（PyG）没有直接的内置API，但可以通过自定义层、边掩码或者封装变换模块来实现，以下是几种实用方案：

方案1：自定义带掩码的GNN层 #

核心思路是在GNN层的forward方法中，仅对无特征节点应用聚合后的特征，有特征节点保留原特征。以GCN为例，你可以基于原生GCNConv改写：

import torch
from torch_geometric.nn import GCNConv
from torch_geometric.utils import add_self_loops, degree

class MaskedGCNConv(GCNConv):
    def forward(self, x, edge_index, mask_no_feat):
        # mask_no_feat是布尔张量，无特征节点位置为True
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        
        # 计算所有节点的聚合特征（和原生GCN逻辑一致）
        x = self.lin(x)
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
        
        aggr_x = torch.sparse.mm(
            torch.sparse_coo_tensor(edge_index, norm, (x.size(0), x.size(0))),
            x
        )
        
        # 仅替换无特征节点的特征，有特征节点保持原输出
        x = torch.where(mask_no_feat.unsqueeze(1), aggr_x, x)
        return x

使用方式 #

先提前生成无特征节点的掩码张量，然后在需要的层使用自定义MaskedGCNConv，其他层用普通GCN：

# 假设你的节点总数为num_nodes，no_feat_idx是无特征节点的索引列表
mask_no_feat = torch.zeros(num_nodes, dtype=torch.bool)
mask_no_feat[no_feat_idx] = True

# 模型定义
class YourModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.gcn1 = GCNConv(in_channels, hidden_channels)
        self.masked_gcn = MaskedGCNConv(hidden_channels, hidden_channels)
        self.gcn3 = GCNConv(hidden_channels, out_channels)
    
    def forward(self, x, edge_index):
        x = self.gcn1(x, edge_index).relu()
        # 仅在第二层执行无特征节点的聚合
        x = self.masked_gcn(x, edge_index, mask_no_feat).relu()
        x = self.gcn3(x, edge_index)
        return x

方案2：生成掩码边索引 #

如果只想让无特征节点聚合有特征节点的信息，可以在特定层生成过滤后的边索引，只保留从有特征节点指向无特征节点的边：

def get_masked_edge_index(edge_index, feat_idx, no_feat_idx):
    # 转换为集合快速判断
    feat_set = set(feat_idx.tolist())
    no_feat_set = set(no_feat_idx.tolist())
    
    # 保留：目标节点是无特征节点，源节点是有特征节点的边
    edge_mask = torch.tensor([
        (col in no_feat_set) and (row in feat_set) 
        for row, col in edge_index.t()
    ], dtype=torch.bool)
    masked_edge_index = edge_index[:, edge_mask]
    
    # 可选：给无特征节点加自环，避免聚合为空
    self_loops = torch.stack([no_feat_idx, no_feat_idx], dim=0)
    masked_edge_index = torch.cat([masked_edge_index, self_loops], dim=1)
    
    return masked_edge_index

使用方式 #

在需要的层传入掩码后的边索引即可：

# 提前定义有特征/无特征节点索引
feat_idx = torch.tensor([0,1,2])
no_feat_idx = torch.tensor([3,4,5])

class YourModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.gcn1 = GCNConv(in_channels, hidden_channels)
        self.gcn2 = GCNConv(hidden_channels, hidden_channels)
        self.gcn3 = GCNConv(hidden_channels, out_channels)
    
    def forward(self, x, edge_index):
        x = self.gcn1(x, edge_index).relu()
        # 第二层使用掩码边索引
        masked_edge_idx = get_masked_edge_index(edge_index, feat_idx, no_feat_idx)
        x = self.gcn2(x, masked_edge_idx).relu()
        # 后续层恢复原边索引
        x = self.gcn3(x, edge_index)
        return x

方案3：封装成可复用的变换模块 #

如果想把聚合逻辑封装成独立模块，方便在特定层前调用，可以写一个类似PyG Transform的类：

class AggregateNoFeatNodes:
    def __init__(self, feat_idx, no_feat_idx):
        self.feat_idx = feat_idx
        self.no_feat_idx = no_feat_idx
    
    def __call__(self, x, edge_index):
        # 筛选出无特征节点与有特征节点之间的边
        edge_mask = torch.isin(edge_index[0], self.feat_idx) & torch.isin(edge_index[1], self.no_feat_idx)
        sub_edge = edge_index[:, edge_mask]
        
        # 计算每个无特征节点的邻居特征均值
        agg_feats = torch.zeros_like(x[self.no_feat_idx])
        for i, node in enumerate(self.no_feat_idx):
            # 获取当前无特征节点的有特征邻居
            neighbors = sub_edge[0][sub_edge[1] == node]
            if len(neighbors) > 0:
                agg_feats[i] = x[neighbors].mean(dim=0)
        
        # 更新无特征节点的特征
        x[self.no_feat_idx] = agg_feats
        return x

使用方式 #

在模型的特定层前调用该变换：

agg_transform = AggregateNoFeatNodes(feat_idx, no_feat_idx)

class YourModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.gcn1 = GCNConv(in_channels, hidden_channels)
        self.gcn2 = GCNConv(hidden_channels, hidden_channels)
        self.gcn3 = GCNConv(hidden_channels, out_channels)
    
    def forward(self, x, edge_index):
        x = self.gcn1(x, edge_index).relu()
        # 在第二层前执行特征聚合
        x = agg_transform(x, edge_index)
        x = self.gcn2(x, edge_index).relu()
        x = self.gcn3(x, edge_index)
        return x

这三种方案都能满足你的需求，其中自定义层的方式最贴合GNN的层流程，边掩码方式能严格控制聚合的数据源，变换模块则更偏向模块化复用，你可以根据自己的代码习惯选择。

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