【基础解读】（PYG）Design of Graph Neural Networks——Heterogeneous Graph Learning

1. Example Graph

在ogbn-mag数据集中，有一个异构图网络，包含1,939,743个节点，分为四种类型：作者、论文、机构和研究领域。它还包含21,111,007条边，边的类型包括：

1. 作者写作特定论文
2. 作者隶属于某个机构
3. 论文引用其他论文
4. 论文涉及特定研究领域

该图的任务是基于图中存储的信息推测每篇论文的发表场所（会议或期刊）。

2. Creating Heterogeneous Graphs

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]

创建一个 torch_geometric.data.HeteroData 类型的数据对象，为每种类型分别定义节点特征张量、边索引张量和边特征张量。
可见——“有一种类型的边，就有一种类型的边索引和边类型的tensor”
此数据对象允许每种类型有不同的特征维度。

可以直接通过 data.{attribute_name}_dict 访问属性分组信息。

model = HeteroGNN(...)

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

现有数据集导入

from torch_geometric.datasets import OGB_MAG

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

输出为：

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 排行榜中的高排名提交中常见的处理方法。

3. Unity Function

单个节点和边单独索引：

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

访问数据对象的元数据：

访问数据对象的元数据，以获取所有现有节点和边类型的信息。

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.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])

4. Heterogeneous Graph Transformations

大多数用于预处理常规图的转换方法同样适用于异构图数据对象。

import torch_geometric.transforms as T

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

1. ToUndirected() 将有向图转换为无向图，通过为每条边添加反向边，使得消息传递可以在所有边的两个方向上进行。如果需要，它会向异构图添加反向边类型。
2. AddSelfLoops() 会为所有类型为 ‘node_type’ 的节点和所有形式为 (‘node_type’, ‘edge_type’, ‘node_type’) 的边类型添加自环边，使得每个节点在消息传递中可能接收到来自自身的消息。
3. NormalizeFeatures() 将所有指定特征（所有类型）标准化，使它们的和为1。

5. Creating Heterogeneous GNNs

标准的消息传递图神经网络（Standard Message Passing GNNs, MP-GNN）不能直接应用于异构图数据，因为不同类型的节点和边特征由于类型差异不能由相同的函数处理。为了解决这个问题，可以为每种边类型单独实现消息和更新函数，MP-GNN在运行时需要迭代边类型字典进行消息计算，并迭代节点类型字典进行节点更新。

PyTorch Geometric 提供了三种创建异构图模型的方法：

1. 使用 torch_geometric.nn.to_hetero() 或 torch_geometric.nn.to_hetero_with_bases() 自动转换同质GNN模型为异构GNN模型；
2. 使用 PyG 的包装器 conv.HeteroConv 为不同类型定义单独的函数；
3. 部署现有或自定义的异构GNN操作符。

Automatically Converting GNN Models

PyTorch Geometric 允许使用内置函数 torch_geometric.nn.to_hetero() 或 torch_geometric.nn.to_hetero_with_bases() 将任何 PyG 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模型复制并修改消息函数，使其能够分别处理每种边类型。

因此，模型现在需要以包含节点和边类型作为键的字典作为输入，而不是在同质图中使用的单一张量。请注意，我们将 in_channels 的元组传递给 SAGEConv，以支持在二分图中的消息传递。

由于不同类型的输入特征数量不同，因此 PyG 可以使用懒加载初始化来初始化异构 GNN 中的参数（如 in_channels 参数为 -1 所示）。这种方式避免了计算和跟踪所有张量尺寸。懒加载初始化适用于所有现有的 PyG 操作符，我们只需要调用一次模型来初始化其参数。

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)

Using the Heterogeneous Convolution Wrapper

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)

Deploy Existing Heterogeneous Operators

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

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)

6. Heterogeneous Graph Samplers

PyG提供了用于采样异构图的多种功能，例如标准的torch_geometric.loader.NeighborLoader类或专门的异构图采样器，如torch_geometric.loader.HGTLoader。这对于大规模异构图上的高效表示学习非常有用，特别是在处理邻居数量过多时。其他采样器，如torch_geometric.loader.ClusterLoader或torch_geometric.loader.GraphSAINTLoader，很快也将支持异构图。所有异构图Loader的输出将是一个HeteroData对象，包含原始数据的子集，主要区别在于采样过程的方式。

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参数，我们进一步指定了要从中采样局部邻居的节点类型和索引。例如，可以选择将数据[‘paper’].train_mask标记为训练节点的所有“paper”节点。

Printing 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包含用于计算128个“paper”节点嵌入的28,187个节点。采样的节点总是按采样顺序进行排序。因此，前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

out[‘paper’][:batch_size]：只要前batch_size个是因为，大于batch_size的节点为前batch_size个节点的邻居节点。

在损失计算中，我们只使用前128个“paper”节点。通过根据batch[‘paper’].batch_size对“paper”标签batch[‘paper’].y和“paper”输出预测out[‘paper’]进行切片，分别表示原始小批量节点的标签和最终输出预测。