如何像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博客