第十九课.Pytorch-geometric扩展

用GraphGym管理实验

介绍

GraphGym 是一个设计和评估图神经网络 (GNN) 的平台，最初在 “Design Space for Graph Neural Networks”论文中提出。 现在已经正式支持 GraphGym 作为 PyG 的一部分。

亮点：
1.GNN的高度模块化的管道（pipeline）：

• Data：数据加载和数据拆分；
• Model：模块化的GNN实现；
• Tasks：节点级别，边级别，图级别的任务；
• Evaluation：Accuracy，ROC，AUC；

2.可扩展的实验管理：

• 并行启动多个GNN实验；
• 自动生成随机种子；

3.灵活的用户自定义：

• 轻松注册自定义模块，比如数据加载器，GNN层，损失函数；

GraphGym 非常适合 GNN 初学者、领域专家和 GNN 研究人员。

场景1：我是图表示学习的初学者，想了解 GNN 的工作原理；

我们可能已经阅读了许多关于 GNN 的激动人心的论文，并尝试编写自己的 GNN 实现。 即使使用原始 PyG，我们仍然必须自己编写必要的模块。 GraphGym 是开始学习标准化 GNN 实现和评估的理想场所；

场景2：可以将 GNN 应用于令人兴奋的应用程序；

我们可能知道有数百种可能的 GNN 模型，众所周知，选择最佳模型非常困难。 更糟糕的是，GraphGym 论文表明，针对不同任务的最佳 GNN 设计差异很大。 GraphGym 提供了一个简单的界面，可以并行尝试数千个 GNN，并了解针对特定任务的最佳设计。 在调查了 1000 万个 GNN 模型-任务组合之后，GraphGym 推荐了一个“首选”GNN 设计空间；

场景 3：对于 GNN 研究人员，想要创新 GNN 模型或提出新的 GNN 任务；

假设提出了一个新的 GNN 层 ExampleConv。 GraphGym 可以帮助证明 ExampleConv 比 GCNConv 更好：当从 1000 万个可能的模型-任务组合中随机抽样时，当其他一切都固定（包括计算成本）时，ExampleConv 能否优于 GCNConv？ 此外，GraphGym 可以帮助进行超参数搜索，并可视化哪些设计选择更好。 总之，GraphGym 可以极大地促进 GNN 研究。

基本用法

要使用 GraphGym，需要从 Github 克隆 PyG，然后切换到 graphgym/ 目录：

git clone https://github.com/pyg-team/pytorch_geometric.git
cd pytorch_geometric/graphgym

然后我们可以在该目录下体验GraohGym的更多用法。

现在GraphGym已经被PyG包含，我们可以在PyG中使用GraphGym：

import torch_geometric.graphgym

小批量处理

小批量的创建对于让深度学习模型扩展到大数据规模是重要的，小批量不是一个一个地处理示例，而是将一组示例分组到统一的表示中，在那里可以有效地并行处理。在图像或语言领域，此过程通常是通过将每个示例重新缩放或填充为一组大小相同的形状来实现的，然后将示例分组到额外的维度中。该维度的长度等于小批量中分组的示例数量，通常称为batch_size；

由于图是最通用的数据结构，可以容纳任意数量的节点和边，因此上述的批量化方法对图不适用；在PyG中，选择另一种方法来实现批量化，在这里，邻接矩阵以对角线方式堆叠（创建一个包含多个孤立子图的巨型图），并且节点特征和目标在节点维度上简单连接。

与其他批处理程序相比，该设计有以下优点：

• 依赖于消息传递模型的GNN算子不需要修改；
• 没有计算或内存上的开销，因为不需要对节点或边缘特征进行填充；

在 torch_geometric.loader.DataLoader 类的帮助下，PyG 会自动将多个图批处理成单个巨型图。在内部，DataLoader 只是一个普通的 PyTorch torch.utils.data.DataLoader，它覆盖了它的 collate() 功能，即示例列表应该如何组合在一起的定义。 因此，所有可以传递给 PyTorch DataLoader 的参数也可以传递给 PyG DataLoader，例如，num_workers；

在最一般的形式中，PyG DataLoader 将根据在当前处理图之前整理的所有图的累积节点数自动增加 edge_index 张量，并将连接 edge_index 张量（形状为 [2, num_edges]）在第二个维度。

但是，有一些特殊用例（如下所述），其中用户希望根据自己的需要修改此行为。 PyG 允许通过覆盖 torch_geometric.data.Data.__inc__() 和 torch_geometric.data.Data.__cat_dim__() 功能来修改底层批处理过程。 在没有任何修改的情况下，这些在 Data 类中定义如下：

def __inc__(self, key, value, *args, **kwargs):
    if 'index' in key or 'face' in key:
        return self.num_nodes
    else:
        return 0

def __cat_dim__(self, key, value, *args, **kwargs):
    if 'index' in key or 'face' in key:
        return 1
    else:
        return 0

我们可以看到 __inc__() 定义了两个连续图属性之间的增量计数，而 __cat_dim__() 定义了应该将同一属性的哪个维度图张量连接在一起。 这两个函数都会为存储在 Data 类中的每个属性调用，并将它们的特定键和值项作为参数传递。

在接下来的内容中，将介绍一些可能绝对需要修改 __inc__() 和 __cat_dim__() 的用例；

Pairs of Graphs

如果想在单个 Data 对象中存储多个图形，例如，对于图形匹配等应用程序，需要确保这些图形的正确批处理行为。比如：考虑存储两个图，Data对象中的一个源（source）图$G_s$一个目标（target）图$G_t$：

from torch_geometric.data import Data

class PairData(Data):
    def __init__(self, edge_index_s=None, x_s=None, edge_index_t=None, x_t=None):
        super().__init__()
        self.edge_index_s = edge_index_s
        self.x_s = x_s
        self.edge_index_t = edge_index_t
        self.x_t = x_t

在这种情况下，edge_index_s 应该增加源图中的节点数，例如 x_s.size(0)，并且 edge_index_t 应该增加目标图中的节点数，例如 x_t.size(0) ：

def __inc__(self, key, value, *args):
    if key == 'edge_index_s':
        return self.x_s.size(0)
    if key == 'edge_index_t':
        return self.x_t.size(0)
    else:
        return super().__inc__(key, value, *args)

我们可以通过设置一个简单的测试脚本来测试我们的 PairData 批处理行为：

from torch_geometric.loader import DataLoader
import torch

edge_index_s = torch.tensor([
     [0, 0, 0, 0],
     [1, 2, 3, 4],
 ])
x_s = torch.randn(5, 16)  # 5 nodes.
edge_index_t = torch.tensor([
     [0, 0, 0],
     [1, 2, 3],
 ])
x_t = torch.randn(4, 16)  # 4 nodes.

data = PairData(edge_index_s, x_s, edge_index_t, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)
# Batch(edge_index_s=[2, 8], x_s=[10, 16], edge_index_t=[2, 6], x_t=[8, 16])

print(batch.edge_index_s)
"""
tensor([[0, 0, 0, 0, 5, 5, 5, 5],
        [1, 2, 3, 4, 6, 7, 8, 9]])
"""

print(batch.edge_index_t)
"""
tensor([[0, 0, 0, 4, 4, 4],
        [1, 2, 3, 5, 6, 7]])
"""

edge_index_s 和 edge_index_t 可以正确地一起批处理，即使在$G_s$和$G_t$使用不同数量的节点时也是如此。 但是，由于 PyG 无法识别 PairData 对象中的实际图，因此缺少批处理属性（将每个节点映射到其各自的图）。 这就是 DataLoader 的 follow_batch 参数发挥作用的地方。 在这里，我们可以指定要维护批次信息的属性：

loader = DataLoader(data_list, batch_size=2, follow_batch=['x_s', 'x_t'])
batch = next(iter(loader))

print(batch) 
#Batch(edge_index_s=[2, 8], x_s=[10, 16], x_s_batch=[10], edge_index_t=[2, 6], x_t=[8, 16], x_t_batch=[8])

print(batch.x_s_batch)
# tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])

print(batch.x_t_batch)
# tensor([0, 0, 0, 0, 1, 1, 1, 1])

follow_batch=['x_s', 'x_t'] 现在成功地分别为节点特征 x_s 和 x_t 创建了名为 x_s_batch 和 x_t_batch 的分配向量。 该信息现在可用于在单个 Batch 对象中的多个图上执行归约操作，例如全局池化。

Bipartite Graphs

二部图的邻接矩阵定义了两种不同节点类型的节点之间的关系。通常，每个节点类型的节点数不需要匹配，其邻接矩阵为 A ∈ { 0 , 1 } N × M , N ≠ M A\in\left\{0,1\right\}^{N\times M},N\neq M A∈{0,1}N×M,N​=M；在二部图的 mini-batching 过程中，edge_index 中边的源节点的增加应该与 edge_index 中边的目标节点不同。 为了实现这一点，考虑两个节点类型之间的二部图，分别具有相应的节点特征 x_s 和 x_t：

from torch_geometric.data import Data
import torch

class BipartiteData(Data):
    def __init__(self, edge_index=None, x_s=None, x_t=None):
        super().__init__()
        self.edge_index = edge_index
        self.x_s = x_s
        self.x_t = x_t

对于二部图中正确的小批量程序，我们需要告诉 PyG 它应该相互独立地增加 edge_index 中边的源节点和目标节点：

def __inc__(self, key, value, *args, **kwargs):
    if key == 'edge_index':
        return torch.tensor([[self.x_s.size(0)], [self.x_t.size(0)]])
    else:
        return super().__inc__(key, value, *args, **kwargs)

这里，edge_index[0]（边的源节点）增加了 x_s.size(0)，而 edge_index[1]（边的目标节点）增加了 x_t.size(0)。 我们可以通过运行一个简单的测试脚本来再次测试我们的实现：

from torch_geometric.loader import DataLoader

edge_index = torch.tensor([
     [0, 0, 1, 1],
     [0, 1, 1, 2],
 ])
x_s = torch.randn(2, 16)  # 2 nodes.
x_t = torch.randn(3, 16)  # 3 nodes.

data = BipartiteData(edge_index, x_s, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)
# Batch(edge_index=[2, 8], x_s=[4, 16], x_t=[6, 16])

print(batch.edge_index)
"""
tensor([[0, 0, 1, 1, 2, 2, 3, 3],
        [0, 1, 1, 2, 3, 4, 4, 5]])
"""

Batching Along New Dimensions

有时，数据对象的属性应该通过获得新的批量维度（如在经典小批量处理中）来进行批量处理，例如，对于图级属性或目标。 具体来说，形状 [num_features] 的属性列表应作为 [num_examples, num_features] 而不是 [num_examples * num_features] 返回。 PyG 通过在 __cat_dim__() 中返回 None 的串联维度来实现这一点：

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import torch

class MyData(Data):
     def __cat_dim__(self, key, value, *args, **kwargs):
         if key == 'foo':
             return None
         else:
             return super().__cat_dim__(key, value, *args, **kwargs)

edge_index = torch.tensor([
   [0, 1, 1, 2],
   [1, 0, 2, 1],
])
foo = torch.randn(16)

data = MyData(edge_index=edge_index, foo=foo)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch) # Batch(edge_index=[2, 8], foo=[2, 16])

根据需要，batch.foo 现在由两个维度描述：批量维度和特征维度。

关于PyG的MessagePassing计算原理

PyG的MessagePassing接口依赖于 gather-scatter （聚集-分散）方案聚合来自相邻节点的消息，比如 message passing 层：$${\text{x}}_i^{\prime }=\sum _{j\in N(i)}MLP({\text{x}}_j-{\text{x}}_i)$$可以被实现为以下伪代码：

from torch_geometric.nn import MessagePassing

x = ...           # Node features of shape [num_nodes, num_features]
edge_index = ...  # Edge indices of shape [2, num_edges]

class MyConv(MessagePassing):
    def __init__(self):
        super().__init__(aggr="add")

    def forward(self, x, edge_index):
        return self.propagate(edge_index, x=x)

    def message(self, x_i, x_j):
        return MLP(x_j - x_i)

在后台，MessagePassing会生成如下伪代码：

from torch_scatter import scatter

x = ...           # Node features of shape [num_nodes, num_features]
edge_index = ...  # Edge indices of shape [2, num_edges]

x_j = x[edge_index[0]]  # Source node features [num_edges, num_features]
x_i = x[edge_index[1]]  # Target node features [num_edges, num_features]

msg = MLP(x_j - x_i)  # Compute message for each edge

# Aggregate messages based on target node indices
out = scatter(msg, edge_index[1], dim=0, dim_size=x.size(0), reduce="sum")

---

torch_scatter.scatter
该方法定义为：

torch_scatter.scatter(
input: torch.Tensor,  # 源张量
index: torch.Tensor,  # 散布元素的索引
dim: int = - 1, # 索引操作的轴
out: Optional[torch.Tensor] = None,  # 目标张量

"""
如果未给出 out，则在维度 dim 处自动创建大小为 dim_size 的输出。 如果未给出 dim_size，则根据 index.max() + 1 作为dim维度的dim_size
"""
dim_size: Optional[int] = None,

reduce: str = 'sum' # 聚合方式
)→ torch.Tensor

其作用是按照索引去聚合分散的值：

举例如下：

from torch_scatter import scatter
import torch

src=torch.randn(1,3,2)
print(src)
#tensor([[[ 0.3450, -0.4956],
#         [-0.4951, -0.6793],
#        [ 0.4820, -1.1948]]])

index = torch.tensor([0,0,1])  
out = scatter(src, index, dim=1, reduce="sum")
# According to index and reduce, we get the sum of 0 and 1 channels' sum.

print(out.size())
#torch.Size([1, 2, 2])

print(out)
#tensor([[[-0.1501, -1.1749],
#        [ 0.4820, -1.1948]]])

---