PyTorch Geometric Introduction by Example 中文

这篇文章是我翻译PyTorch Geometric官方文档中的一篇教程，水平有限，要是有什么错误，希望能多加担待。欢迎大家指出我的错误与不足指出。

通过例子介绍

我们通过一些独立的例子简短地介绍下PyTorch Geometric的一些重要概念。在其内核中，PyTorch Geometric提供了以下几个主要特征:

• 图数据处理
• 通用基准数据集
• 小批次
• 数据转换
• 图学习方法

图数据处理

图用于构建对象(节点)之间成对关系(边)的模型。在PyTorch Geometric中，单个图通过一个torch_geometric.data.Data的实例化经行描述。这个图默认拥有以下几种属性：

• data.x：形如[num_nodes, num_node_features]的节点特征矩阵
• data.edge_index：用形状为[2, num_edges]并且类型为torch.long的COO格式描述图连通性
• data.edge_attr：形如[num_edges, num_edge_features]的边特征矩阵
• data.y：针对训练的标签(可能拥有任意形状)，例如，形为[num_nodes, *]的节点层次标签或形为[1, *]的图层次标签
• data.pos：形如[num_nodes, num_dimensions]的节点位置矩阵

上述属性没有一个是必须的。事实上，Data对象甚至不受这些属性的制约。例如，我们可以将从3D网格中的三角形连通性保存在形如[3, num_faces]并且类型为torch.long的张量(tensor)中，并通过data.face进行扩充。

---

！Note

Pytorch和torchvision定义一个像图片与标签的元组的样本(这句话的意思大概是PyTorch和torchvision定义一个样本时，就是将图片与标签组成一个元组)。我们在PyTorch Geometric中忽略了这种表示方法，从而使各种数据结构可以通过一种简单并且易于理解的方法进行定义。

---

我们展示一个具有三个节点和四条边的无权重的无向图的简单例子。每个节点仅包含一个特征：

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-624DA0n2-1607772987746)(D:\Only_bases\Only_photo\PyTorch_geometric\graph.svg)]

import torch
from torch_geometric.data import Data
edgge_index = torch.tensor([[0, 1, 1, 2],
                            [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)
>>>Data(edge_index=[2, 4], x=[3, 1])

值得注意的是edge_index，即，这个张量定义了所有边的源和目的节点，而不是索引元组的列表。如果你想通过元组方式构建你的边，你应转置边并在将其送入data的构造函数之前调用contiguous方法：

import torch
from torch_geometric.data import Data
edge_index = torch.tensor([[0, 1],
                           [1, 0],
                           [1, 2],
                           [2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
data = Data(x=x, edge_index.t().contiguous())
>>>Data(edge_index=[2, 4], x=[3, 1])

尽管这个图只有两条边，我们也是需要定义四组索引元组，用于说明每个边的方向。

---

！Note

你可以在任何时刻打印你的数据对象，然后获得关于它的属性和形状的简短信息

---

除了一些普通的python对象的功能，Data还提供了很多有效的功能，例如：

print(data.keys)
>>> ['x', 'edge_index']

print(data['x'])
>>> tensor([[-1.0],
            [0.0],
            [1.0]])

for key, item in data:
    print("{} found in data".format(key))
>>> x found in data
>>> edge_index found in data

'edge_attr' in data
>>> False

data.num_nodes
>>> 3

data.num_edges
>>> 4

data.num_node_features
>>> 1

data.contains_isolated_nodes()
>>> False

data.contains_self_loops()
>>> False

data.is_directed()
>>> False

# Transfer data object to GPU.
device = torch.device('cuda')
data = data.to(device)

你可以在torch_geometric.data.Data中获取全部方法的列表。

通用基准数据集

PyTorch Geometric包含了大量的通用基准数据集，例如，所有类行星数据集(Cora, Citeseer, Pubmed)，从 http://graphkernels.cs.tu-dortmund.de中和他们的清理版本所获得的全部图分类数据集，QM7和QM9数据集，像FAUST、ModelNet10/40、ShapeNet这样的3D网/点云数据集。

初始化一个数据集是直接了当的。初始化一个数据集会自动的下载raw文件并且将其处理为Data格式的数据。例如，加载ENZYMES数据集(共6个分类的600个图)：

from torch_geometric.datasets import TUDataset

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
>>> ENZYMES(600)

len(dataset)
>>> 600

dataset.num_classes
>>> 6

dataset.num_node_features
>>> 3

现在，我们可以访问数据集中所有的600个图：

data = dataset[0]
>>> Data(edge_index=[2, 168], x=[37, 3], y=[1])

data.is_undirected()
>>> True

我们可以看出，在数据集的第一个图中拥有37个节点，每个节点包含3个特征。这个图含有168/2=84个无向边，并且这个图属于一个确定的类。除此之外，这个数据对象恰好持有一个图层次标签。

我们可以使用slices、long、byte张量来切片数据集。例如，创建一个90/10的训练/测试切片：

train_dataset = dataset[:540]
>>> ENZYMES(540)

test_dataset = dataset[540:]
>>> ENZYMES(60)

如果你不确定数据集是否在你切片前已经乱序过了，你可以运行如下程序进行乱序：

dataset = dataset.shuffle()
>>> ENZYMES(600)

等价的做法：

perm = torch.randperm(len(dataset))
dataset = dataset[perm]
>> ENZYMES(600)

让我们尝试一下其他的数据集。让我们下载Cora——半监督图节点分类标准基准数据集：

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()

len(dataset)
>>> 1

dataset.num_classes
>>> 7

dataset.num_node_features
>>> 1433

在这里，数据集仅包含一个无向的引文图：

data = dataset[0]
>>> Data(edge_index=[2, 10556], test_mask=[2708],
         train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])

data.is_undirected()
>>> True

data.train_mask.sum().item()
>>> 140

data.val_mask.sum().item()
>>> 500

data.test_mask.sum().item()
>>> 1000

此刻，data对象持有每个节点的标签，并且包含额外的属性：train_mask，val_mask，test_mask：

• train_mask表明针对哪些节点经行训练(140个节点)
• val_mask表明哪些节点用于验证
• test_mask表明针对哪些节点进行测试(1000个节点)

小批次

神经网络通常是通过小批次方式进行训练。PyTorch Geometric通过创建稀疏块对角线邻接矩阵（由edge_index定义），并在节点维度中合并特征矩阵和目标矩阵，从而在小型批处理上实现并行化。
$$A=\left[\begin{array}{ccc}{A}_1& & \\ & \ddots & \\ & & A_n\end{array}\right],X=\left[\begin{array}{c}{X}_1\\ \cdots \\ X_n\end{array}\right],Y=\left[\begin{array}{c}{Y}_1\\ \cdots \\ Y_n\end{array}\right]$$

PyTorch Geometric包含自己的torch_geometric.data.DataLoader，它已经负责此串联过程。让我们通过一个例子进行学习：

from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    batch
    >>> Batch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])

    batch.num_graphs
    >>> 32

torch_geometric.data.Batch内置于torch_geometric.data.Data中，并且包含一个额外的属性batch。

batch是一个列向量，它将每个节点映射到该批次中的相应图：
$$batch={\left[\begin{array}{ccccccccc}0& \cdots & 0& 1& \cdots & n-2& n-1& \cdots & n-1\end{array}\right]}^T$$
例如，你可以使用它来分别地在每个图形的节点维度中平均节点特征：

from torch_scatter import scatter_mean
from torch_geometric.datasets import TUDataset
from torch_geometric.data import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for data in loader:
    data
    >>> Batch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])

    data.num_graphs
    >>> 32

    x = scatter_mean(data.x, data.batch, dim=0)
    x.size()
    >>> torch.Size([32, 21])

从这你可以学习到更多关于PyTorch Geometric中内部的批处理的过程，例如，如何修改它的做法。有关离散操作的文档，请向感兴趣的读者阅读torch_scatter文档。

数据转换

变换是torchvision中变换图像和执行增强的常见方式。PyTorch Geometric拥有它自己的转换方式，它期望一个Data对象作为输入并且返回一个转换的新的Data对象。可以使用torch_geometric.transforms.Compose将变换链接在一起，并在将处理后的数据集保存到磁盘之前(pre_transform)或访问数据集中的图之前应用变换(transform)。

让我们看一个示例，其中我们对ShapeNet数据集（包含17,000个3D形状点云和来自16个形状类别的每个点标签）应用变换。

from torch_geometric.datasets import ShapeNet

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'])

dataset[0]
>>> Data(pos=[2518, 3], y=[2518])

我们可以通过转换从点云生成最近的邻居图，从而将点云数据集转换为图数据集：

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

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6))

dataset[0]
>>> Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])

---

! Note

在将数据存储到磁盘之前，我们使用pre_transform来变换数据(为了加快加载速度)。请注意，下一次初始化数据集时，即使您不传递任何变换，它也会包含图形的边。

---

此外，我们可以使用transform参数随机扩展Data对象，例如，用一个更小的数字对每个节点位置经行变换：

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

dataset = ShapeNet(root='/tmp/ShapeNet', categories=['Airplane'],
                    pre_transform=T.KNNGraph(k=6),
                    transform=T.RandomTranslate(0.01))

dataset[0]
>>> Data(edge_index=[2, 15108], pos=[2518, 3], y=[2518])

您可以在torch_geometric.transforms中找到所有已实现转换的完整列表。

图学习方法

经过对PyTorch Geometric中的数据处理、数据集、加载和变换的学习，现在是时候部署我们的第一个图神经网络了。

我们将使用一个简单的GCN层并将实验复制到Cora引用数据集上。有关GCN的高层次解释，请查看其博客文章。

首先，我们需要加载Cora数据集：

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')
>>> Cora()

请注意，我们不需要使用transforms或dataloader。现在，我们实现一个两层GCN:

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

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

构造函数定义了两个GCNConv层，它们在我们网络的前向传播中被调用。请注意，非线性功能并未集成在conv调用中，因此需要在之后应用(在PyTorch Geometric中，所有运算符都保持一致)。在这里，我们选择使用ReLU作为介于两者之间的中间非线性，最后输出整个类数的softmax分布。让我们在训练节点上训练这个模型200个轮：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

最后，在测试节点上评估我们的模型：

model.eval()
_, pred = model(data).max(dim=1)
correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('Accuracy: {:.4f}'.format(acc))
>>> Accuracy: 0.8150

这就是实现您的第一个图神经网络所需要的。了解更多关于图卷积和池化的最简单方法是研究examples /目录中的示例并浏览torch_geometric.nn。