pytorch geometric（PYG） - NeighborSampler

          如何像graphsage中对mini-batch的节点进行邻居采样并训练模型，使得大规模全连接图的GNN模型训练成为可能，pyg是通过torch_geometric.loader.NeighborSampler实现的；

        只要卷积层支持二分图，就可以与NeighborSampler结合使用；

1 参数介绍  

        NeighborSampler：它允许在完全批量训练不可行的情况下，对大规模图上的gnn进行小批量训练；

        给定一个具有:math: ' L '层的GNN和一个特定的小批节点:obj: ' node_idx '，我们想要计算嵌入，这个模块迭代采样邻居，并构建二分图来模拟GNN的实际计算流程；

        更具体地说，:obj: ' sizes '表示我们希望在每个层中的每个节点采样多少邻居；

        该模块然后接受这些:obj: ' size '，并迭代采样:obj: ' sizes[l] '，每个节点涉及层:obj: ' l '。在下一层，对已经遇到的节点的并集重复采样；

        然后以反向模式返回实际的计算图，这意味着我们将消息从较大的节点集传递到较小的节点集，直到到达我们最初想要计算嵌入的节点集。

        因此，由:class: ' NeighborSampler '返回的一个项保存当前:obj: ' batch_size '，所有参与计算的节点的id: obj: ' n_id '，以及通过元组: obj: ' (edge_index, e_id, size) '的二分图对象列表，其中:obj: ' edge_index '表示源节点和目标节点之间的二分图边，obj: ' e_id '表示完整图中原始边的id，和:obj: ' size '保存二分图的形状。

        对于每个二分图，目标节点也包括在源节点列表的开头，以便可以轻松地应用跳过连接或添加自循环。

        二分图：二分图又称作二部图，是图论中的一种特殊模型。 设G=(V,E)是一个无向图，如果顶点V可分割为两个互不相交的子集(A,B)，并且图中的每条边（i，j）所关联的两个顶点i和j分别属于这两个不同的顶点集(i in A,j in B)，则称图G为一个二分图。

        区别二分图，关键是看点集是否能分成两个独立的点集：（所有回路的长度均为偶数）

 

def __init__(self, edge_index: Union[Tensor, SparseTensor],
                 sizes: List[int], node_idx: Optional[Tensor] = None,
                 num_nodes: Optional[int] = None, return_e_id: bool = True,
                 transform: Callable = None, **kwargs):

        edge_index (Tensor or SparseTensor)：图的边信息，可以是Tensor，也可以是SparseTensor；

        sizes ([int])：每一层需要采样的邻居数目，如果是-1的话，选取所有的邻居；

        node_idx (LongTensor, optional)：提供需要被采样节点的信息，比如模型训练的时候，只给出数据集train中的节点。在预测的时候，使用None，考虑所有的节点。

        num_nodes: Optional[int] = None：图中节点的数目，可选参数。

        return_e_id: bool = True：当设为False的时候，不会返回partite子图的边在原图中的IDs。

        transform

        **kwargs：NeighborSampler是torch.utils.data.DataLoader的子类，所以父类DataLoader的参数NeighborSampler都可以使用，比如:batch_size, shuffle, num_workers 

2 核心想法

        给定mini-batch的节点和图卷积的层数L，以及每一层所需要的采样邻居的数目 sizes，依次从第1层到第L层，对每一层进行邻居采样，并返回二部图，sizes是一个长度为L的list，包含每一层所需要采样的邻居个数；具体实施过程可以去看B站视频：16. 4.3_GraphSAGE代码_哔哩哔哩_bilibili

        每一层采样返回结果： (edge_index, e_id, size);

                edge_index是采样得到的bipartite子图中source节点到target节点的边;

                e_id是edge_index的边在原始大图中的IDs;

                size就是bipartite子图的shape;

        L 层采样完成后，返回结果：(batch_size, n_id, adjs)

                batch_size就是mini-batch的节点数目；

                n_id：L层采样中遇到的所有的节点的list，其中target节点在list最前端；

                adjs：第L层到第1层采样结果的list；

               edge_index：采样得到的bipartite子图中source节点到target节点的边；

                                  e_id：edge_index的边在原始大图中的IDs；

                                  size：bipartite子图的shape；

3 定义train, dev, test_loader

        定义方式：

train_loader = NeighborSampler(edge_index=edge_index, node_idx=user_index, 
sizes=[-1], batch_size=len(user_index))
# edge_index (Tensor or SparseTensor)：图的边信息，可以是Tensor，也可以是SparseTensor；
 
#node_idx (LongTensor, optional)：提供需要被采样节点的信息，比如模型训练的时候，只给出数据集#train中的节点。在预测的时候，使用None，考虑所有的节点
 
 
 
# 或者
 
train_loader = NeighborSampler(data.edge_index, node_idx=data.train_mask,
                               sizes=[25, 10], batch_size=1024, shuffle=True,
                               num_workers=12)
#node_idx=data.train_mask指定只对训练集的节点进行邻居采样;
# sizes=[25, 10]指明了这是一个两层的卷积，第一层卷积采样邻居数目25，第二层卷积采样邻居数目10;
# batch_size=1024指定了mini-batch的节点数目，每次只对1024个节点进行采样

        NeighborSampler是在CPU中完成的，所以返回的结果都在CPU上。如果用GPU训练模型，要记得将loader的结果放到GPU上;

        train_loader每次返回一个batch_size节点邻居采样的结果;

4 模型训练

代码展示：

def train(epoch):
    model.train()
    
    total_loss = total_correct = 0
    for batch_size, n_id, adjs in train_loader:
        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
        adjs = [adj.to(device) for adj in adjs]
 
        optimizer.zero_grad()
        # x[n_id]是所有相关节点的特征 x[n_id]相当于做了一次映射，x[n_id]中第i行就是adjs中i节点的特征
        # adjs是包含了所有bipartite子图边信息的list
        # model(x[n_id], adjs)传入了所有bipartite子图的节点特征和边信息
        out = model(x[n_id], adjs)
        loss = F.nll_loss(out, y[n_id[:batch_size]])
        loss.backward()
        optimizer.step()
 
        total_loss += float(loss)
        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
 
    loss = total_loss / len(train_loader)
    approx_acc = total_correct / int(data.train_mask.sum())
    
    return loss, approx_acc

 forward函数：

        实现了从第L层到第1层采样得到的bipartite子图的卷积；

class SAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
        super(SAGE, self).__init__()
        self.num_layers = num_layers
        ...
 
    # x是一个tuple: (x_source, x_target)
    def forward(self, x, adjs):
        for i, (edge_index, _, size) in enumerate(adjs): # 对于所有节点，利用一阶邻居更新embedding
            x_target = x[:size[1]]  # Target nodes are always placed first.
            # 实现了对一层bipartite图的卷积
            x = self.convs[i]((x, x_target), edge_index)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)
        return x.log_softmax(dim=-1)

bipartite图的size是(num_of_source_nodes, num_of_target_nodes)，因此对每一层的bipartite图都有 x_target = x[:size[1]]；

        不是取所有的n阶邻居，计算一次得到节点最终的嵌入；（这种方法很需要内存，但是GPU有时内存不够）

        而是每次只取所有的一阶邻居，但是进行n次迭代，然后得到中心节点的嵌入；

5 NeighborSampler工作原理&具体实例

        networkx支持创建简单无向图、有向图和多重图（multigraph）；内置许多标准的图论算法，节点可为任意数据；支持任意的边值维度，功能丰富，简单易用。

        首先使用networkx构建一张图：

import networkx as nx
graph = nx.Graph()
graph.add_edges_from([(0,1), (1,2), (1,3), (2,3), (3,4), (4,2)])
nx.draw_kamada_kawai(graph, with_labels=True)

        将其转为PYG中的Data格式

from torch_geometric.data.data import Data
from torch_geometric.utils import from_networkx
 
data = from_networkx(graph)
data.edge_index
>>> tensor([[0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4],
            [1, 0, 2, 3, 1, 3, 4, 1, 2, 4, 2, 3]])

         采样邻居数小于邻居数

from torch_geometric.data import NeighborSampler
# sizes ([int])：每一层需要采样的邻居数目，如果是-1的话，选取所有的邻居
loader = NeighborSampler(edge_index=data.edge_index, sizes=[2], node_idx=torch.tensor([2]), batch_size=1)
next(iter(loader))
 
# batch_size
# n_id：L层采样中遇到的所有的节点的list，其中target节点在list最前端
# 第L层到第1层采样结果的list
    # 采样得到的bipartite子图中source节点到target节点的边
>>> (1, tensor([2, 3, 1]),
     EdgeIndex(edge_index=tensor([[1, 2],[0, 0]]), e_id=tensor([8, 2]), size=(3, 1)))

        以上代码对2号节点进行邻居采样，n_id: tensor([2,3,1]), 是采取到的所有节点;

        target节点在最前面，是2号节点，3, 1是采样到的邻居;

        edge_index=tensor([[1, 2], [0, 0]])是采样得到的bipartite子图;

        n_id中的index对应edge_index中的数值，edge_index[1]中是target节点       

复制自： （作者讲的很清楚，我在这里做记录使用，希望大家去看原作者）pytorch geometric教程四 利用NeighorSampler实现节点维度的mini-batch + GraphSAGE样例_每天都想躺平的大喵的博客-CSDN博客