概述

PyTorch Geometric (PyG) 是一个用于几何深度学习的 python 框架,可以在图形、点云和流形等不规则结构上实现图神经网络 (GNN)。它提供专用的 CUDA 内核来有效处理稀疏数据,并提供小批量处理程序以适应不同大小的数据。PyTorch Geometric 能够在 CPU 和 GPU 上运行,其优化的吞吐量适用于高度稀疏和不规则的数据。

什么是PyG(PyTorch Geometric)?

PyTorch Geometric (PyG) 是一个 python 库,用于对不规则结构(如图形、点云和流形)进行深度学习。它建立在流行的 PyTorch 框架之上,并在可变大小的稀疏和不规则数据上提供 GPU 加速计算的优化吞吐量。

PyG 库包含来自已发表论文的各种几何深度学习方法,可用于执行节点分类、图形分类或链接预测等任务。此外,它还包含用于操作许多小型或单个巨型图形的小型批量加载器、多 GPU 支持、DataPipe 支持、通过 Quiver 进行分布式图形学习,以及用于任意图形以及 3D 网格或点云的有用转换。

它是一个直观而强大的框架,使研究人员能够轻松地为与结构化数据相关的各种应用实现图神经网络。

在图形结构数据上试用 PyTorch Geometric 的原因是什么?

  • 优化了高度稀疏和不规则数据的吞吐量:
    PyTorch Geometric 为稀疏数据提供专用的 CUDA 内核和不同大小的迷你批量处理程序,使用户在处理图形结构数据时能够充分利用其 GPU 资源。
  • 支持 CPU 和 GPU 计算:
    PyTorch Geometric 能够在 CPU 和 GPU 上运行,无论硬件限制如何,各种用户都可以使用它。
  • 易于使用的工具:
    PyTorch Geometric 包含易于使用的小型批量加载器以及有用的转换,使研究人员能够以最小的努力快速实现与图形结构相关的学习任务。
  • 全面的方法库:
    该库包含从已发表论文中进行几何深度学习的各种方法,使用户能够轻松地将最新研究应用于他们的数据。
  • 灵活的库:
    PyTorch Geometric 是一个灵活的库,可用于节点分类、图形分类或链接预测等任务。

要求和安装

在继续安装之前,有一些要求。他们是:

  • Python 3.6 或更高版本。
  • PyTorch 1.5 或更高版本(建议使用 1.7)。
  • CUDA 10+(用于 GPU 支持)。

安装

步骤 - 1:通过运行(CPU 版本)安装 PyTorch

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu

如果您需要 GPU 支持,请确保您的系统在运行的 PC 中安装了兼容的 CUDA 版本

nvcc --version

您将获得已安装的 cuda 驱动程序的版本。

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Feb_14_21:12:58_PST_2021
Cuda compilation tools, release 11.2, V11.2.152
Build cuda_11.2.r11.2/compiler.29618528_0

然后使用以下命令安装支持 CUDA 的 PyTorch

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

更多自定义安装请参考 PyTorch 官方文档。

步骤 - 2:使用以下命令(CPU 版本)安装 PyG 库和其他支持库:

pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.13.0+cpu.html

对于 GPU 版本,您应该安装兼容的 CUDA 驱动程序和 PyTorch GPU。然后使用以下命令安装 PyG GPU

pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.13.0+cu116.html

步骤 - 3:使用以下命令安装其他用于可视化的软件包,例如 matplotlib 和 networks:

pip install networkx[default]

创建和训练 GNN 模型

我们可以使用图形数据集,例如 Zachary's Karate Club,用于节点分类任务。该数据集有 34 个节点代表俱乐部成员,78 个链接代表他们之间的互动。节点被标记为属于两个派系之一。我们可以使用训练集来构建图神经网络模型,然后使用该模型来预测测试集中节点的派系。

Zachary 空手道俱乐部的标准 78 边缘网络数据集在互联网上公开提供。数据可以汇总为整数对列表。每个整数代表一个空手道俱乐部成员,一对表示两个成员互动。数据集总结如下。

[2 1]
[3 1] [3 2]
[4 1] [4 2] [4 3]
[5 1]
[6 1]
[7 1] [7 5] [7 6]
[8 1] [8 2] [8 3] [8 4]
[9 1] [9 3]
[10 3]
[11 1] [11 5] [11 6]
[12 1]
[13 1] [13 4]
[14 1] [14 2] [14 3] [14 4]
[17 6] [17 7]
[18 1] [18 2]
[20 1] [20 2]
[22 1] [22 2]
[26 24] [26 25]
[28 3] [28 24] [28 25]
[29 3]
[30 24] [30 27]
[31 2] [31 9]
[32 1] [32 25] [32 26] [32 29]
[33 3] [33 9] [33 15] [33 16] [33 19] [33 21] [33 23] [33 24] [33 30] [33 31] [33 32]
[34 9] [34 10] [34 14] [34 15] [34 16] [34 19] [34 20] [34 21] [34 23] [34 24] [34 27] [34 28] [34 29] [34 30] [34 31] [34 32] [34 33]

Zachary 的空手道俱乐部图形数据集的简单可视化如下所示:

 首先,我们从 PyG 导入必要的模块,例如 torch、GCNConv 和 KarateClub 数据集。

import os
import networkx as nx
import matplotlib.pyplot as plt
import torch
from torch.nn import Linear
from torch_geometric.nn import GCNConv

PyTorch Geometric 通过 torch_geometric.datasets 子包提供对此数据集的轻松访问:

from torch_geometric.datasets import KarateClub

dataset = KarateClub()
print(f'Dataset: {dataset}:')
print('======================')
print(f'Number of graphs: {len(dataset)}')
print(f'Number of features: {dataset.num_features}')
print(f'Number of classes: {dataset.num_classes}')

初始化空手道俱乐部数据集后,我们首先可以检查它的一些属性。例如,我们可以看到这个数据集正好包含一个图形,并且该数据集中的每个节点都被分配了一个 34 维特征向量(它唯一地描述了空手道俱乐部的成员)。此外,该图正好包含 4 个类,它们代表每个节点所属的社区。

收集有关图形的信息

现在让我们更详细地看一下底层图:

data = dataset[0]  # Get the first graph object.

print(data)
print('==============================================================')

# Gather some statistics about the graph.
print(f'Number of nodes: {data.num_nodes}')
print(f'Number of edges: {data.num_edges}')
print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')
print(f'Number of training nodes: {data.train_mask.sum()}')
print(f'Training node label rate: {int(data.train_mask.sum()) / data.num_nodes:.2f}')
print(f'Has isolated nodes: {data.has_isolated_nodes()}')
print(f'Has self-loops: {data.has_self_loops()}')
print(f'Is undirected: {data.is_undirected()}')

PyTorch Geometric 中的每个图形都由一个 Data 对象表示,该对象包含描述其图形表示的所有信息。我们可以随时通过 print(data) 打印数据对象,以接收其属性及其形状的摘要:

Data(edge_index=[2, 156], x=[34, 34], y=[34], train_mask=[34])

我们可以看到这个数据对象包含 4 个属性:

  1. edge_index 属性包含有关图形连通性的信息,即每个边的源节点索引和目标节点索引的元组。PyG 进一步指
  2. 节点特征为 x(34 个节点中的每一个都被分配了一个 34 个 dim 的特征向量),并且
  3. 节点标记为 y(每个节点正好分配给一个类)。
  4. 还有一个称为 train_mask 的附加属性,它描述了我们已经知道哪些节点的社区分配。总的来说,我们只知道 4 个节点(每个社区一个)的真值标签,任务是推断其余节点的社区分配。

现在让我们更详细地检查edge_index属性:

edge_index = data.edge_index
print(edge_index.t())

通过打印edge_index,我们可以理解 PyG 如何在内部表示图连通性。我们可以看到,对于每条边,edge_index包含一个由两个节点索引组成的元组,其中第一个值描述源节点的节点索引,第二个值描述边目标节点的节点索引。

这种表示形式称为 COO 格式(坐标格式),通常用于表示稀疏矩阵。

 

使用 NetworkX 可视化图形

我们可以通过将图形转换为网络库格式来进一步可视化图形,除了图形操作功能外,它还实现了强大的可视化工具:

from torch_geometric.utils import to_networkx

G = to_networkx(data, to_undirected=True)
visualize_graph(G, color=data.y)

实现图神经网络

在了解了 PyG 的数据处理之后,是时候实现我们的第一个图神经网络了。PyG 通过 GCNConv 实现该层,可以通过传入节点特征表示 x 和 COO 图连通性表示edge_index来执行。

有了这个,我们就可以通过在 torch.nn.Module 类中定义我们的网络架构来创建我们的第一个图神经网络:

import torch
from torch.nn import Linear
from torch_geometric.nn import GCNConv


class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        torch.manual_seed(1234)
        self.conv1 = GCNConv(dataset.num_features, 4)
        self.conv2 = GCNConv(4, 4)
        self.conv3 = GCNConv(4, 2)
        self.classifier = Linear(2, dataset.num_classes)

    def forward(self, x, edge_index):
        h = self.conv1(x, edge_index)
        h = h.tanh()
        h = self.conv2(h, edge_index)
        h = h.tanh()
        h = self.conv3(h, edge_index)
        h = h.tanh()  # Final GNN embedding space.

        # Apply a final (linear) classifier.
        out = self.classifier(h)

        return out, h

model = GCN()

在这里,我们首先在 __init__ 中初始化所有构建块,并在前向定义网络的计算流。

我们首先定义并堆叠三个图卷积层,它们对应于聚合每个节点周围的 3 跳邻域信息(所有节点最多 3 个“跳”距离)。

此外,GCNConv 层将节点特征维数降低到22,即34→4→4→234→4→4→2.每个 GCNConv 层都通过 tanh 非线性进行增强。

之后,我们应用一个线性变换torch.nn.Linear,它充当分类器,将我们的节点映射到4个类/社区中的1个。

我们返回最终分类器的输出以及 GNN 生成的最终节点嵌入。

我们继续通过 GCN() 初始化我们的最终模型,并打印我们的模型生成其所有使用的子模块的摘要。

嵌入可视化

让我们看一下我们的 GNN 生成的节点嵌入。在这里,我们将初始节点特征 x 和edge_index到模型的图形连通性信息传入模型,并可视化其二维嵌入。

model = GCN()

_, h = model(data.x, data.edge_index)
print(f'Embedding shape: {list(h.shape)}')

visualize_embedding(h, color=data.y)

值得注意的是,甚至在训练模型的权重之前,该模型就产生了与图的社区结构非常相似的节点嵌入。相同颜色(社区)的节点已经在嵌入空间中紧密地聚集在一起,尽管我们模型的权重是完全随机初始化的,到目前为止我们还没有进行任何训练!由此得出的结论是,GNN引入了强烈的归纳偏置,导致输入图中彼此靠近的节点具有相似的嵌入。

空手道俱乐部网络培训

训练我们的模型与任何其他 PyTorch 模型非常相似。

除了定义我们的网络架构之外,我们还定义了一个损失准则并初始化一个随机梯度优化器。

之后,我们执行多轮优化,其中每一轮都由前向和后向传递组成,以计算模型参数相对于正向传递得出的损失的梯度。

model = GCN()
criterion = torch.nn.CrossEntropyLoss()  # Define loss criterion.
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # Define optimizer.

def train(data):
    optimizer.zero_grad()  # Clear gradients.
    out, h = model(data.x, data.edge_index)  # Perform a single forward pass.
    loss = criterion(out[data.train_mask], data.y[data.train_mask])  # Compute the loss solely based on the training nodes.
    loss.backward()  # Derive gradients.
    optimizer.step()  # Update parameters based on gradients.
    return loss, h

for epoch in range(401):
    loss, h = train(data)
    if epoch % 10 == 0:
        visualize_embedding(h, color=data.y, epoch=epoch, loss=loss)
        time.sleep(0.3)

正如人们所看到的,我们的 3 层 GCN 模型设法线性地分离社区并正确分类大多数节点。

此外,我们只用了几行代码就完成了这一切,这要归功于 PyTorch Geometric 库,它帮助我们完成了数据处理和 GNN 实现。

结论

本文向我们展示了什么是 GNN 以及如何使用 PyG 构建 GNN。我们学到了以下几点:

  • 图神经网络 (GNN) 因其能够捕获图结构化数据的结构并将其用于学习任务而变得越来越流行。
  • PyTorch Geometric (PyG) 是一个 python 库,使用户能够以最小的努力轻松实现 GNN。
  • 它提供了对稀疏和不规则数据的优化吞吐量,支持 CPU 和 GPU 计算、易于使用的工具(如小型批量加载器)和有用的转换,以及来自已发表论文的综合方法库。
  • 因此,它是一个直观而强大的框架,可用于与结构化数据相关的各种应用程序。
Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐