玩转大型异构图，Graphcore扩展IPU PyTorch Geometric支持

大型异构图在现实世界的许多数据集中都十分常见。

在这种情况下，“大型”是指单个图中具有数千个节点和边，“异构”指的是图中的节点代表着不同类型的实体，而边代表着这些实体类型之间的各种关系。

例如，社交网络可以用图的形式建模，节点代表用户，节点之间的边表示友谊关系。异构图可以容纳其他节点类型，如用户发布的帖子、他们所属的群组以及参加的活动。

一个完整的社交网络可以变得非常庞大。它可以包含许多用户、帖子等等，并包含它们之间的所有关系。

在最新发布的Graphcore（拟未） Poplar SDK 3.3中，我们扩展了PyTorch Geometric IPU支持，使用户可以利用Graphcore IPU来加速这类应用。

在本文中，我们将简要介绍支持在IPU上使用大型异构图的最新功能。我们还有许多关于该主题的教程，这些教程可以在Paperspace Gradient Notebooks上运行。除此之外，我们还提供了一个示例来展示在大型异构图上使用GNN进行欺诈检测

^[1]

。

https://ipu.dev/uN5BYV

https://ipu.dev/0xyjM7

https://ipu.dev/XB1fAX

IPU上的大型图

随着图尺寸的增加，对于处理器来说，全批训练所需的存储容量会在某个点变得有点太大了——即使对拥有业界领先的片上SRAM的IPU来说也是如此。

全批处理中，模型的输入是整个图，每次迭代都涉及所有的节点和边。

解决该问题的方法是从大型图中采样，形成较小的迷你批作为模型的输入。

一种常用的方法是GraphSAGE相邻采样^[2]，即从节点中抽取一个迷你批来计算一个表示，然后随机选取这些节点的相邻节点，以及相邻节点的相邻节点等等。通过这种方式，可以形成目标节点的良好表示，同时具有可扩展性。

使用PyTorch Geometric是一种非常简单的方法。NeighborLoader对象提供一个数据加载器，用于生成迷你批样本。

IPU采用提前编译，这意味着整个图必须是静态的，包括输入。这样就能有效布局存储和通信，并在编译期间进行某些优化。

PyTorch Geometric提供的NeighborLoader生成的迷你批并不是固定尺寸的，因此在Graphcore最新的Poplar SDK 3.3 PopTorch Geometric包中，我们提供了一个FixedSizeNeighborLoader对象，它封装了PyG的NeighborLoader，并额外固定了输出的迷你批尺寸。它操作简单，可以轻松替换：

- from torch_geometric.loader import NeighborLoader

+ from poptorch_geometric.neighbor_loader import FixedSizeNeighborLoader

- neighbor_loader = NeighborLoader(data, num_neighbors=[-1, -1], input_nodes=("movie", data["movie"].train_mask), batch_size=128)

+ neighbor_loader = FixedSizeNeighborLoader(data, num_neighbors=[-1, -1], input_nodes=("movie", data["movie"].train_mask), batch_size=128)

此外，我们还提供了FixedSizeClusterLoader，它相当于PyG簇加载器，具有固定大小的迷你批输出。

更多有关如何在IPU上启用大型图的信息请参阅我们的教程^[3]。

IPU上的异构图

如前所述，现实世界中的许多图都是异构的，包含多种类型的节点以及它们之间的多种类型关系。PyTorch Geometric已经为使用异构图提供了强大的支持，支持用户以灵活简洁的方式构建模型。

更多信息，请参见PyG文档

^[4]

。

在IPU上使用PyG的异构功能非常简单，您只需要使用现有功能，并将您的模块包装在另一个包含损失函数的模块中即可。

heterogeneous_model = HeterogeneousGNN()

class ModelWithLoss(torch.nn.Module):

def __init__(self, model):

super().__init__()

self.model = model

def forward(self, x_dict, edge_index_dict, target=None, train_mask=None):

out = self.model(x_dict, edge_index_dict)

if self.training:

target = torch.where(train_mask, target, -100)

loss = F.cross_entropy(out["target_nodes"], target)

return out, loss

return out

Graphcore最新的Poplar SDK 3.3 PopTorch Geometric包中的全部现有数据加载器现都支持异构图数据，使其可以支持创建固定尺寸的异构迷你批。

为实现这一点，我们可以通过提供大量的节点和边，使迷你批达到固定尺寸，我们也可以通过对每个节点和边的类型设置不同的值，从而减少迷你批中包含的用于填充的节点和边。

它还包括辅助功能，可自动从现有的PyG动态数据加载器中获取所需数量的节点和边，以填充迷你批。

from poptorch_geometric import FixedSizeOptions

from poptorch_geometric.neighbor_loader import FixedSizeNeighborLoader

neighbor_loader = NeighborLoader(dataset, num_neighbors=[-1, -1] , input_nodes=("movie", data["movie"].train_mask), batch_size=128)

fixed_size_options = FixedSizeOptions.from_loader(neighbor_loader)

print(fixed_size_options)

FixedSizeOptions(

num_nodes={

'movie': 64,

'director': 6,

'actor': 16}

(At least one node reserved for padding),

num_edges={

('movie', 'to', 'director'): 22,

('movie', 'to', 'actor'): 56,

('director', 'to', 'movie'): 6,

('actor', 'to', 'movie'): 16}

(At least one edge reserved for padding),

num_graphs=2 (At least one graph reserved for padding))

neighbor_loader = FixedSizeNeighborLoader(dataset, num_neighbors=[-1, -1] , batch_size=128, input_nodes=("movie", data["movie"].train_mask), fixed_size_options=fixed_size_options)

我们编写了一份详细的教程^[5]帮助您了解更多。

结论

您可以利用上述对使用PyG采样方法的额外支持，以及对如何在IPU上使用PyG异构功能的指导，简单直接地在Graphcore IPU上使用大型图进行异构图学习。

您可以查看我们的教程和Paperspace Gradient Notebooks示例。这些教程和示例支持Paperspace的六小时免费IPU试用。