使用torch_geometric.loader.NeighborLoader创建子图时如何保留全部边？

先明确根因：PyTorch Geometric（PyG）的NeighborLoader目前不支持同时开启disjoint=True和directed=False，而且默认采样逻辑会过滤掉部分边，导致边特征丢失。结合你的节点分类需求，给你三个可行的替代方案：

方案1：用ClusterLoader做无重叠子图划分（最贴合disjoint需求） #

如果你的图可以做聚类划分，ClusterLoader能直接把大图切成不相交的子图，而且会完整保留每个子图里的所有边和特征，完美适配大尺度图的节点分类：

from torch_geometric.loader import ClusterLoader
from torch_geometric.transforms import RandomNodeSplit

# 先划分训练/验证/测试节点
data = RandomNodeSplit()(data)
# 用METIS算法把图分成100个聚类（数量根据你的图规模调整）
loader = ClusterLoader(
    data,
    num_clusters=100,
    shuffle=True,
    batch_size=10,  # 一次加载10个聚类子图
)

# 训练循环示例
for batch in loader:
    x, edge_index, edge_attr = batch.x, batch.edge_index, batch.edge_attr
    out = model(x, edge_index, edge_attr)
    loss = criterion(out[batch.train_mask], batch.y[batch.train_mask])

这个方案完全避开了NeighborLoader的限制，还能保证子图边特征不丢失，对无向图友好。

方案2：自定义非随机全邻居采样（保留disjoint=False但不丢边） #

如果必须用邻居采样模式，别开disjoint=True，改用disjoint=False，同时设置num_neighbors=[-1]实现非随机的全邻居采样，这样Loader会自动保留子图内的所有边特征：

from torch_geometric.loader import NeighborLoader

# num_neighbors=[-1]表示取每个节点的所有一阶邻居，实现非随机全采样
loader = NeighborLoader(
    data,
    num_neighbors=[-1],
    batch_size=256,
    shuffle=False,  # 关闭打乱，保证采样固定
    directed=False,  # disjoint=False时支持无向图
    input_nodes=data.train_mask,
)

# 训练时直接用batch里的edge_attr就行，不会丢
for batch in loader:
    out = model(batch.x, batch.edge_index, batch.edge_attr)
    loss = criterion(out, batch.y[batch.batch])

注意：这个模式下子图会有重叠节点，训练时正常更新梯度就行，PyG会自动处理重复节点的梯度累加。

方案3：用GraphSAINTLoader做节点采样（专为大图优化） #

如果你的任务接受随机采样，GraphSAINTLoader是PyG专为大尺度图设计的Loader，能保证采样的子图包含所有关联边，边特征完全保留，而且天然支持无向图：

from torch_geometric.loader import GraphSAINTNodeSampler

loader = GraphSAINTNodeSampler(
    data,
    batch_size=256,
    num_steps=10,
    sample_coverage=100,
    shuffle=True,
)

for batch in loader:
    out = model(batch.x, batch.edge_index, batch.edge_attr)
    loss = criterion(out, batch.y)

这个Loader的采样效率很高，适合超大规模图的训练。

补充说明 #

你遇到的Undirected subgraphs not yet supported报错，是因为PyG至今（包括最新版）还没在disjoint=True模式下实现无向图的子图划分逻辑，所以这俩参数不能同时设。上面三个方案都绕开了这个限制，同时满足你保留边特征、小批量训练的需求。

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