在Graphcore拟未IPU上使用PyTorch Geometric的实用攻略

2023/04/05

在Graphcore拟未IPU上使用PyTorch Geometric的实用攻略

作者:

在这篇文章中，我们将展示如何在IPU上轻松开始使用PyG。我们将讲解加速IPU上训练的端到端示例，并提供一些提示和技巧，来帮助您熟悉PyG在IPU上的使用。最后，我们将介绍PopTorch Geometric数据加载器，它能够使您在GNN工作负载中最大限度地利用IPU。

同时也欢迎查看我们的博客，可以了解更多关于IPU如何以及为什么适合您的GNN工作负载。

您可以直接在Paperspace上开始免费学习，例如通过探索一系列的入门教程、一些更深入的应用示例或基准测试来展示IPU在消息传递操作上的出色表现。

在IPU上免费试用PyG

移植到IPU：基础知识

要在IPU上使用PyTorch Geometric运行工作负载，模型需要以PopTorch为目标。PopTorch是一套针对IPU的扩展，允许您在IPU上运行PyTorch原生模型。它的设计要求尽可能少地改变原生PyTorch，但也有一些需要注意的差异，我们将在本节的其他部分探讨这些差异。

让我们从使用PyG的典型模型开始：


import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        torch.manual_seed(1234)
        self.conv = GCNConv(in_channels, out_channels, add_self_loops=False)

    def forward(self, x, edge_index, edge_weight=None):
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv(x, edge_index, edge_weight).relu()
        return x


model = GCN(dataset.num_features, dataset.num_classes)
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

print("Training on CPU.")

for epoch in range(1, 6):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index, data.edge_attr)
    loss = F.cross_entropy(out, data.y)
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch}, Loss: {loss}")

您可能已经熟悉了上述代码块中的概念。我们创建了一个简单的模型，由一个从PyG导入的GCNConv层组成，选择了一个优化器并构建了一个训练循环，我们在其中计算损失并进行优化器更新。

现在我们可以通过最小的更改，让这个模型可以在IPU上运行：

在上面的视频中，导入PopTorch后，我们已经：

将损失函数移到模型的forward方法中。
将模型包裹在poptorch.trainingModel中，因为我们计划在这个代码片段中运行训练。如果我们计划运行推理，则会用poptorch.inferenceModel来包裹模型。
使用专为IPU设计的PopTorch优化器。
删除了手动调用后向传递和优化器的步骤，这两个步骤都是由PopTorch自动处理的。

下面展示的是准备在IPU上运行的最终代码：


import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

import poptorch

class GCN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        torch.manual_seed(1234)
        self.conv = GCNConv(in_channels, out_channels, add_self_loops=False)

    def forward(self, x, edge_index, y, edge_weight=None):
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv(x, edge_index, edge_weight).relu()

        if self.training:
            loss = F.cross_entropy(x, y)
            return x, loss
        return x


model = GCN(dataset.num_features, dataset.num_classes)
model.train()
optimizer = poptorch.optim.Adam(model.parameters(), lr=0.001)
poptorch_model = poptorch.trainingModel(model, optimizer=optimizer)

print("Training on IPU.")

for epoch in range(1, 6):
    output, loss = poptorch_model(
        data.x, data.edge_index, data.y, edge_weight=data.edge_attr
    )
    print(f"Epoch: {epoch}, Loss: {loss}")

然后就可以了！通过这些简单的改变，您可以成功地在IPU上运行PyG模型。现在，我们准备使用一个端到端的示例进行更深入的研究。

端到端示例：基于GCN的简单GNN，用于节点分类

让我们在一个示例中应用上述概念，我们将使用一个简单的模型对Cora数据集的节点进行分类，该模型由几个GCN层组成。Cora数据集是一个引文网络，其中一个节点代表一个文档，如果两个文档之间有引文，则存在边缘。

首先，加载数据集：


from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T

transform = T.Compose([T.NormalizeFeatures(), T.AddSelfLoops()])

dataset = Planetoid(root=dataset_directory, name="Cora", transform=transform)
data = dataset[0]  # Access the citation graph as Data object

PyG提供了很多有用的层，让您只需几行代码就能构建GNN。对于这项任务，我们将使用GCNConv层，因为GCN是最常用的GNN运算符之一。

在定义这个模型之前，我们首先应该了解，IPU上的PyTorch支持依赖于通过PopTorch进行的AOT（ahead-of-time）编译。编译器对图进行静态分析：为了优化IPU上的评估，需要输入张量形状的先验知识。我们在当前和后面的示例中所做的众多更改都是为了确保模型能够被成功编译。

我们现在准备根据一些观察结果来定义模型：

GCNConv自循环被禁用。这是因为在按照上面的代码片段加载数据集时，我们已经将自循环作为一个transform添加到数据集中，因此我们可以在层中关闭它们。同时，为了确保为IPU编译的图是静态的，我们也需要进行这一更改。
PopTorch要求损失函数是模型的一部分，所以我们将把它移到forward方法中。
我们将把非训练节点的标签设置为-100，以便在损失函数中忽略它们：此步骤使用torch.where()，确保y是已知大小，满足IPU上静态大小输入的要求。


import torch
from torch_geometric.nn import GCNConv
import torch.nn.functional as F


class GCN(torch.nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels, 16, add_self_loops=False)
        self.conv2 = GCNConv(16, out_channels, add_self_loops=False)

    def forward(self, x, edge_index, y=None, train_mask=None):
        x = self.conv1(x, edge_index).relu()
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index).relu()
        x = F.log_softmax(x, dim=1)

        if self.training:
            y = torch.where(train_mask, y, -100)
            loss = F.nll_loss(x, y)
            return x, loss
        return x

然后将模型包裹在PopTorch trainingModel()中，我们将使用PopTorch优化器（在本例中为Adam）。使用PopTorch特定的优化器可以帮助提高速度和存储的利用率：这些优化器使用起来非常简单，因为它们是现有PyTorch优化器的直接替代品。


model = GCN()
model.train()
optimizer = poptorch.optim.Adam(model.parameters(), lr=0.001)
poptorch_model = poptorch.trainingModel(model, optimizer=optimizer)

现在我们可以运行训练，并确保减少损失：


for epoch in range(1, 101):
    out, loss = poptorch_model(
        data.x, data.edge_index, y=data.y, edge_weight=data.edge_attr, train_mask=data.train_mask
    )
    print(f"Epoch: {epoch}, Loss: {loss}")

这就是在IPU上训练我们的模型所需的全部工作！下一步是推理。

运行推理

为了在PopTorch中运行推理，我们将使用训练好的模型，将其设置为eval模式，并包裹在PopTorch inferenceModel()中。


model.eval()
poptorch_inf_model = poptorch.inferenceModel(model)

由于我们不再进行训练，因此不需要定义损失，所以我们可以在模型本身不做任何掩蔽。相反，可以事后在CPU上进行掩蔽。首先，我们从模型中得到预测结果：


data = next(iter(dataset))
logits = poptorch_inf_model(data.x, data.edge_index)
pred = logits.argmax(dim=1)

现在我们可以在CPU上计算验证精度了。这些相对便宜的操作是可行的，没有什么可以阻止我们将它们放入我们的模型定义中，并在IPU上进行计算。


correct_results = pred[data.val_mask] == data.y[data.val_mask]
accuracy = int(correct_results.sum()) / int(data.val_mask.sum())
print(f"Validation accuracy: {accuracy:.2%}")

这样我们就获得了训练模型的验证精度。

这是一个非常简单的例子，我们可以在其中进行完整的批量更新，也就是说，我们可以在单批次中处理完整的图。随着输入的图大小的增加，这种情况可能并不总是存在。在下一节内容中，我们将展示不能或可能不需要进行完整批量更新的用例：这时数据加载就会发挥作用。让我们来看看如何充分利用IPU上的数据加载器。

利用PopTorch Geometric的高性能数据加载器

任何模型在训练和推理中的一个关键部分都是准备数据、批量处理并将其输入模型。

如果您正在使用GNN，您的输入数据可能是一组可以放在一起进行批处理的小型图，也可能是一个需要采样或完整批量更新的大型图。

在上一节中，我们强调了IPU需要固定大小的输入，这意味着需要对输入张量的形状的先验知识进行了解。有其他不同的方法来实现这一点，所使用的方法取决于我们正在处理的输入图数据集的类型：

如果要处理一个由许多小型图组成的数据集，可以通过数据加载器对输入图进行分批处理，并对生成的批量进行填充。利用PopTorch Geometric中提供的最新工具，使用FixedSizeDataLoader就可以非常简单地实现这一点。您可以查看我们关于使用填充进行小型图批处理的教程，了解详细的操作方法。在特定的使用情况下，这种方法可能会导致大量的填充：我们在关于使用打包进行小型图批处理的专门教程中介绍了一种更有效的批处理策略，即打包。
如果我们处理的是单个大型图的数据集，我们可以从中取样，然后填充样本以获得静态形状。在PopTorch Geometric中，我们提供了一个工具来生成固定大小的集群，以用于大型图工作负载，即FixedSizeClusterLoader。您可以参考Cluster CGN的示例来了解大型图的用例。
如果您的图是中等大小，可以进行完整的批量更新，那么固定大小的要求就变得微不足道了，因为只有一个批次，每个批次的大小都是一样的。我们在上面的端到端示例中看到了这种简单的情况

使用PopTorch Geometric中提供的任何数据加载器，就像上面提到的那些，还可以解锁一些有助于加速模型的功能，例如：