异构图学习

大量的现实世界数据集以异构图的形式存储，这促使在中为它们引入专门的功能。 例如，推荐领域中的大多数图，如社交图，都是异构的，因为它们存储了关于不同类型实体及其不同类型关系的信息。 本教程介绍了如何将异构图映射到，以及如何将它们用作图神经网络模型的输入。

异构图附带有不同类型的节点和边信息。 因此，由于类型和维度的差异，单个节点或边特征张量无法保存整个图的所有节点或边特征。 相反，需要分别为节点和边指定一组类型，每种类型都有其自己的数据张量。 由于不同的数据结构，消息传递公式也相应改变，允许根据节点或边类型计算消息和更新函数。

示例图表

作为一个指导性的例子，我们来看一下来自 数据集套件的异构ogbn-mag网络：

../_images/hg_example.svg

给定的异构图有1,939,743个节点，分为四种节点类型：作者、论文、机构和研究领域。 它还有21,111,007条边，这些边也属于四种类型之一：

  • writes: 一位作者 撰写 了一篇特定的论文

  • 隶属于: 作者隶属于特定机构

  • 引用: 一篇论文引用另一篇论文

  • 有主题: 一篇论文有一个主题，属于特定研究领域

该图的任务是根据图中存储的信息推断每篇论文的发表地点（会议或期刊）。

创建异构图

首先，我们可以创建一个类型为torch_geometric.data.HeteroData的数据对象，我们为每种类型分别定义节点特征张量、边索引张量和边特征张量：

from torch_geometric.data import HeteroData

data = HeteroData()

data['paper'].x = ... # [num_papers, num_features_paper]
data['author'].x = ... # [num_authors, num_features_author]
data['institution'].x = ... # [num_institutions, num_features_institution]
data['field_of_study'].x = ... # [num_field, num_features_field]

data['paper', 'cites', 'paper'].edge_index = ... # [2, num_edges_cites]
data['author', 'writes', 'paper'].edge_index = ... # [2, num_edges_writes]
data['author', 'affiliated_with', 'institution'].edge_index = ... # [2, num_edges_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_index = ... # [2, num_edges_topic]

data['paper', 'cites', 'paper'].edge_attr = ... # [num_edges_cites, num_features_cites]
data['author', 'writes', 'paper'].edge_attr = ... # [num_edges_writes, num_features_writes]
data['author', 'affiliated_with', 'institution'].edge_attr = ... # [num_edges_affiliated, num_features_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_attr = ... # [num_edges_topic, num_features_topic]

节点或边的张量将在首次访问时自动创建，并通过字符串键进行索引。 节点类型由单个字符串标识，而边类型通过使用字符串的三元组(source_node_type, edge_type, destination_node_type)来标识：边类型标识符和边类型可以存在的两个节点类型。 因此，数据对象允许每种类型具有不同的特征维度。

包含按属性名称而不是按节点或边类型分组的异构信息的字典可以直接通过 data.{attribute_name}_dict 访问，并作为后续GNN模型的输入：

model = HeteroGNN(...)

output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

如果数据集存在于Pytorch Geometric数据集列表中，它可以直接导入并使用。 特别是，它将被下载到root并自动处理。

from torch_geometric.datasets import OGB_MAG

dataset = OGB_MAG(root='./data', preprocess='metapath2vec')
data = dataset[0]

data 对象可以打印出来进行验证。

HeteroData(
  paper={
    x=[736389, 128],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389]
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  field_of_study={ x=[59965, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)

注意

原始的ogbn-mag网络仅提供“论文”节点的特征。 在OGB_MAG中，我们提供了下载其处理版本的选项，其中结构特征（从"metapath2vec"或"TransE"获得）被添加到无特征的节点中，正如在OGB排行榜上排名靠前的提交中常见的那样。

实用函数

torch_geometric.data.HeteroData 类提供了许多有用的实用函数来修改和分析给定的图。

例如，单个节点或边存储可以单独索引：

paper_node_data = data['paper']
cites_edge_data = data['paper', 'cites', 'paper']

如果边类型可以通过源节点类型和目标节点类型的唯一对或边类型唯一标识，以下操作同样适用：

cites_edge_data = data['paper', 'paper']
cites_edge_data = data['cites']

我们可以添加新的节点类型或张量并删除它们：

data['paper'].year = ...    # Setting a new paper attribute
del data['field_of_study']  # Deleting 'field_of_study' node type
del data['has_topic']       # Deleting 'has_topic' edge type

我们可以访问data对象的元数据，其中包含所有现有节点和边类型的信息：

node_types, edge_types = data.metadata()
print(node_types)
['paper', 'author', 'institution']
print(edge_types)
[('paper', 'cites', 'paper'),
('author', 'writes', 'paper'),
('author', 'affiliated_with', 'institution')]

data 对象可以像往常一样在设备之间传输：

data = data.to('cuda:0')
data = data.cpu()

我们还可以访问其他辅助函数来分析给定的图

data.has_isolated_nodes()
data.has_self_loops()
data.is_undirected()

并且可以通过to_homogeneous()将其转换为同质的“类型化”图，该函数能够在不同类型的特征维度匹配时保持特征：

homogeneous_data = data.to_homogeneous()
print(homogeneous_data)
Data(x=[1879778, 128], edge_index=[2, 13605929], edge_type=[13605929])

这里，homogeneous_data.edge_type 表示一个边级别的向量，它以整数形式保存每条边的类型。

异构图转换

大多数用于预处理常规图的变换同样适用于异构图data对象。

import torch_geometric.transforms as T

data = T.ToUndirected()(data)
data = T.AddSelfLoops()(data)
data = T.NormalizeFeatures()(data)

在这里，ToUndirected() 将一个有向图转换为（ 表示的）无向图，通过为图中的所有边添加反向边。 因此，未来的消息传递将在所有边的两个方向上进行。 如果有必要，该函数可能会向异构图添加反向边类型。

对于所有类型为'node_type'的节点和所有形式为('node_type', 'edge_type', 'node_type')的现有边类型，函数AddSelfLoops()将添加自循环边。 因此，在消息传递过程中，每个节点可能会从自身接收一条或多条（每种适当的边类型一条）消息。

转换 NormalizeFeatures() 的工作方式与同质情况相同，并将所有指定类型的所有特征归一化为总和为一。

创建异构GNNs

标准消息传递图神经网络（MP-GNNs）不能直接应用于异构图数据，因为不同类型的节点和边特征由于特征类型的差异，无法通过相同的函数进行处理。 一个自然的解决方法是分别为每种边类型实现消息和更新函数。 在运行时，MP-GNN算法需要在消息计算期间迭代边类型字典，并在节点更新期间迭代节点类型字典。

为了避免不必要的运行时开销，并使创建异构MP-GNN尽可能简单，Pytorch Geometric为用户提供了三种方法来在异构图数据上创建模型：

  1. 通过使用torch_geometric.nn.to_hetero()或torch_geometric.nn.to_hetero_with_bases()，自动将同质GNN模型转换为异质GNN模型。

  2. 使用包装器torch_geometric.nn.conv.HeteroConv为不同类型的异构卷积定义单独的函数

  3. 部署现有的（或编写您自己的）异构GNN操作符

以下详细介绍每个选项。

自动转换GNN模型

Pytorch Geometric 允许使用内置函数 torch_geometric.nn.to_hetero() 或 torch_geometric.nn.to_hetero_with_bases() 自动将任何 GNN 模型转换为适用于异构图输入的模型。 以下 示例 展示了如何应用它：

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x


model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

该过程采用现有的GNN模型，并复制消息函数以分别处理每种边类型，如下图所示。

../_images/to_hetero.svg

因此，模型现在期望输入参数是包含节点和边类型作为键的字典，而不是在同质图中使用的单一张量。 请注意，我们传递了一个in_channels的元组给SAGEConv，以便在二分图中进行消息传递。

注意

由于输入特征的数量以及张量的大小在不同类型之间有所不同， 可以利用延迟初始化来初始化异构 GNN 中的参数（如 -1 作为 in_channels 参数所示）。 这使我们能够避免计算和跟踪计算图中所有张量的大小。 所有现有的 操作符都支持延迟初始化。 我们可以通过调用一次模型来初始化其参数：

with torch.no_grad():  # Initialize lazy modules.
    out = model(data.x_dict, data.edge_index_dict)

无论是 to_hetero() 还是 to_hetero_with_bases()，在将同构架构自动转换为异构架构方面都非常灵活，例如，支持开箱即用地应用跳跃连接、跳跃知识或其他技术。 例如，以下代码展示了如何实现一个带有可学习跳跃连接的异构图注意力网络：

from torch_geometric.nn import GATConv, Linear, to_hetero

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x


model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

请注意，我们通过add_self_loops=False参数禁用了自环的创建。 这是因为在二分图中，自环的概念并不明确（对于具有不同源节点和目标节点类型的边类型，消息传递的概念不明确），我们可能会错误地将边[(0, 0), (1, 1), ...]添加到二分图中。 为了保留中心节点的信息，我们因此通过conv(x, edge_index) + lin(x)使用可学习的跳跃连接，这将执行从源节点特征到目标节点特征的基于注意力的消息传递，然后将其输出与现有的目标节点特征相加。

之后，创建的模型可以像往常一样进行训练：

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x_dict, data.edge_index_dict)
    mask = data['paper'].train_mask
    loss = F.cross_entropy(out['paper'][mask], data['paper'].y[mask])
    loss.backward()
    optimizer.step()
    return float(loss)

使用异构卷积包装器

异构卷积包装器 torch_geometric.nn.conv.HeteroConv 允许定义自定义的异构消息和更新函数，以从头开始构建任意用于异构图的MP-GNN。 虽然自动转换器 to_hetero() 对所有边类型使用相同的操作符，但该包装器允许为不同的边类型定义不同的操作符。 在这里，HeteroConv 接受一个子模块字典作为输入，每个子模块对应图数据中的一种边类型。 以下 示例 展示了如何应用它。

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                ('paper', 'cites', 'paper'): GCNConv(-1, hidden_channels),
                ('author', 'writes', 'paper'): SAGEConv((-1, -1), hidden_channels),
                ('paper', 'rev_writes', 'author'): GATConv((-1, -1), hidden_channels, add_self_loops=False),
            }, aggr='sum')
            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        return self.lin(x_dict['author'])

model = HeteroGNN(hidden_channels=64, out_channels=dataset.num_classes,
                  num_layers=2)

我们可以通过调用模型一次来初始化它（有关延迟初始化的更多详细信息，请参见这里）

with torch.no_grad():  # Initialize lazy modules.
     out = model(data.x_dict, data.edge_index_dict)

并运行标准训练程序，如此处所述。

部署现有的异构操作符

提供了专门为异构图设计的操作符（例如，torch_geometric.nn.conv.HGTConv）。 这些操作符可以直接用于构建异构图神经网络模型，如下面的示例所示：

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import HGTConv, Linear


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
        super().__init__()

        self.lin_dict = torch.nn.ModuleDict()
        for node_type in data.node_types:
            self.lin_dict[node_type] = Linear(-1, hidden_channels)

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HGTConv(hidden_channels, hidden_channels, data.metadata(),
                           num_heads, group='sum')
            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for node_type, x in x_dict.items():
            x_dict[node_type] = self.lin_dict[node_type](x).relu_()

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)

        return self.lin(x_dict['author'])

model = HGT(hidden_channels=64, out_channels=dataset.num_classes,
            num_heads=2, num_layers=2)

我们可以通过调用模型一次来初始化它（有关延迟初始化的更多详细信息，请参见这里）。

with torch.no_grad():  # Initialize lazy modules.
     out = model(data.x_dict, data.edge_index_dict)

并按照这里概述的标准训练程序运行。

异构图采样器

提供了多种功能用于采样异构图，即在标准的 torch_geometric.loader.NeighborLoader 类中或专门的异构图采样器如 torch_geometric.loader.HGTLoader 中。 这对于在大型异构图上进行高效的表示学习特别有用，因为处理所有邻居在计算上过于昂贵。 其他采样器如 torch_geometric.loader.ClusterLoader 或 torch_geometric.loader.GraphSAINTLoader 的异构图支持将很快添加。 总的来说，所有异构图加载器都将生成一个 HeteroData 对象作为输出，该对象包含原始数据的子集，并且主要区别在于它们的采样过程的工作方式。 因此，只需进行最少的代码更改即可将训练过程从 全批次训练 转换为小批次训练。

使用NeighborLoader进行邻居采样的工作方式如下示例所示：

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.loader import NeighborLoader

transform = T.ToUndirected()  # Add reverse edge types.
data = OGB_MAG(root='./data', preprocess='metapath2vec', transform=transform)[0]

train_loader = NeighborLoader(
    data,
    # Sample 15 neighbors for each node and each edge type for 2 iterations:
    num_neighbors=[15] * 2,
    # Use a batch size of 128 for sampling training nodes of type "paper":
    batch_size=128,
    input_nodes=('paper', data['paper'].train_mask),
)

batch = next(iter(train_loader))

值得注意的是，NeighborLoader 适用于同质图和异质图。 在操作异质图时，可以对单个边类型的采样邻居数量进行更细粒度的控制，但这不是必须的，例如：

num_neighbors = {key: [15] * 2 for key in data.edge_types}

使用input_nodes参数，我们进一步指定了要从中采样本地邻域的节点的类型和索引，即根据data['paper'].train_mask标记为训练节点的所有“paper”节点。

打印 batch 然后产生以下输出：

HeteroData(
  paper={
    x=[20799, 256],
    y=[20799],
    train_mask=[20799],
    val_mask=[20799],
    test_mask=[20799],
    batch_size=128
  },
  author={ x=[4419, 128] },
  institution={ x=[302, 128] },
  field_of_study={ x=[2605, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 0] },
  (author, writes, paper)={ edge_index=[2, 5927] },
  (paper, cites, paper)={ edge_index=[2, 11829] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 10573] },
  (institution, rev_affiliated_with, author)={ edge_index=[2, 829] },
  (paper, rev_writes, author)={ edge_index=[2, 5512] },
  (field_of_study, rev_has_topic, paper)={ edge_index=[2, 10499] }
)

因此，batch 总共包含28,187个节点，这些节点用于计算128个“paper”节点的嵌入。 采样的节点总是按照它们被采样的顺序进行排序。 因此，前 batch['paper'].batch_size 个节点代表原始小批量节点的集合，这使得通过切片轻松获得最终的输出嵌入。

在迷你批次模式下训练我们的异质GNN模型与在全批次模式下训练类似，除了我们现在迭代由train_loader生成的迷你批次，并基于单个迷你批次优化模型参数：

def train():
    model.train()

    total_examples = total_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        batch = batch.to('cuda:0')
        batch_size = batch['paper'].batch_size
        out = model(batch.x_dict, batch.edge_index_dict)
        loss = F.cross_entropy(out['paper'][:batch_size],
                               batch['paper'].y[:batch_size])
        loss.backward()
        optimizer.step()

        total_examples += batch_size
        total_loss += float(loss) * batch_size

    return total_loss / total_examples

重要的是，我们在损失计算过程中仅使用前128个“paper”节点。 我们通过基于batch['paper'].batch_size来切片“paper”标签batch['paper'].y和“paper”输出预测out['paper']来实现这一点，分别代表原始小批量节点的标签和最终输出预测。