第十八课.Pytorch-geometric入门(三)

异构图学习 heterogeneous graph learning

首先了解一些基础概念：

• 同构图：节点和边的类型均只有一种，比如社交网络中只存在一种节点类型（用户节点）和一种边类型（用户-用户的连边）；
• 异构图：节点的类型或者边的类型有多种，比如论文引用网络，节点有两种类型（作者节点和文章节点），边的类型有作者-作者间的共同创作关系，作者-文章间的从属关系，文章-文章间的引用关系；
• 属性图：图的节点上存在初始属性，可以用作后续节点的特征；
• 动态图：图中的节点或者边都是随着时间变化的，可能增加也可能减少，每个时刻用一个图表示；
• 关系图：图表示了一种节点之间的隐含关系，比如知识图谱；

大量真实世界的数据集存储为异构图，这促使在 Pytorch Geometric (PyG) 中为它们引入专门的功能。 例如，推荐领域的大多数图，都是异构的，因为它们存储有关不同类型实体及其不同类型关系的信息。 下面介绍异构图如何映射到 PyG 以及如何将它们作为图神经网络模型的输入；

异构图带有附加到节点和边的不同类型的信息。 因此，由于类型和维度的差异，单个节点或边特征张量不能包含整个图的所有节点或边特征。 相反，需要分别为节点和边指定一组类型，每个类型都有自己的数据张量。 由于不同的数据结构的影响，消息传递公式相应地改变，允许以节点或边类型为条件的message和update函数的计算；

---

目录中标注（示例）的部分代表：该部分内容对PyG及其组件有版本要求，因为PyG的版本问题会存在一些api不可用的现象；

要使用到一些特殊的api，我们需要安装未发布的PyG（2.0.0），我们可以在原有PyG的环境下安装，安装过程如下：

conda install git
pip install git+https://github.com/pyg-team/pytorch_geometric.git

# 未来PyG2.0.0发布后，我们可以用pip install package==2.0.0的方式安装

PyG对于异构图的处理比较不方便，建议使用DGL（Deep Graph Library）；

---

示例图

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

给定的异构图有 1,939,743 个节点，分为作者（Author）、论文（Paper）、机构（Institution）和研究领域（Field of Study）四种节点类型。 它还有 21,111,007 条边，它们也是四种类型之一：

• writes：作者写了一篇特定的论文；
• affiliated with：作者隶属于特定机构；
• cites：一篇论文引用了另一篇论文；
• has topic：一篇论文有一个特定研究领域的主题；

该图的任务是根据存储在图中的信息推断每篇论文所在的场所venue（会议或者期刊），在ogbn-mag中总共有349个不同的venue，使任务成为349类的分类问题；

创建异构图

首先，我们可以创建一个类型为 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['author', 'has_topic', 'institution'].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]

可以打印数据对象进行验证：

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

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

数据对象可以在GPU与CPU设备上传输：

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 表示一个边级别的向量，它将每条边的边类型保存为一个整数。

异构图的转换

大多数用于预处理常规图数据对象的 transformations 也适用于异构图数据对象。

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() 将添加自循环边。 因此，在消息传递期间，每个节点可能会从其自身接收一个或多个消息；

Normalize Features() 的工作方式类似于同构情况，并将所有指定的特征（所有类型的）归一化。

创建异构GNN

标准消息传递 GNN（MP-GNN）不能轻易应用于异构图数据，因为由于特征类型的差异，不同类型的节点和边缘特征不能由相同的函数处理。 规避这种情况的一种自然方法是为每个边缘类型单独实现消息和更新功能。 在运行时，MP-GNN 算法需要在消息计算期间迭代边缘类型字典，在节点更新期间迭代节点类型字典。

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

• Automatically Converting GNN Models：通过使用 torch_geometric.nn.to_hetero() 或 torch_geometric.nn.to_hetero_with_bases() 自动将同构 GNN 模型转换为异构 GNN 模型；
• Using the Heterogenous Convolution Wrapper：使用 torch_geometric.nn.conv.HeteroConv 为异构卷积定义不同类型的单个函数；
• Deploy Existing Heterogenous Operators：部署现有（或编写自己的）异构 GNN 算子；

Automatically Converting GNN Models（示例）

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

from torch_geometric.nn import SAGEConv, to_hetero

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模型，复制消息函数以单独处理每种边类型；

因此，模型的输入是以节点和边类型构成的字典，不再是同构图中的简单张量；

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) # -1表示延迟初始化
        self.lin1 = Linear(-1, hidden_channels) # -1表示延迟初始化
        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')

然后，可以照常训练创建好的模型：

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[mask], data['paper'].y[mask])
    loss.backward()
    optimizer.step()
    return float(loss)

补充内容：同构图下的GATConv

GATConv来自 “GRAPH ATTENTION NETWORKS”，是一种在GCN消息聚合时增加注意力机制的操作；

以常规图为计算对象，GATConv层的输入为$\text{h}=[h_1,h_2,...,h_N],h_i\in R^F$，其中$N$是节点的数量，$F$是输入特征的维数，输出为${\text{h}}^{\prime }=[h_1^{\prime },h_2^{\prime },...,h_N^{\prime }],h_i^{\prime }\in R^{F^{\prime }}$，$F^{\prime }$是输出特征的维数；

为了获得足够的表达能力将输入特征转换为更高级别的特征，至少需要一种可学习的线性变换。 为此，作为初始步骤，共享线性变换，由权重矩阵$W\in R^{F^{\prime }\times F}$变换；

在计算某个自注意力分布时，我们借助全连接网络实现。${\alpha }_{ij}$表示邻居节点$j$到当前节点$i$的注意力：$${\alpha }_{ij}=\frac{exp(LeakyReLU(a^T[W{h}_i\oplus W{h}_j]))}{{\sum }_{k\in N(i)\cup i}exp(LeakyReLU(a^T[W{h}_i\oplus W{h}_k]))}$$其中，$a\in R^{2F^{\prime }}$是该自注意力全连接网络的参数。

然后将注意力融合到消息传递过程中得到GCN的输出：$$h_i^{\prime }={\alpha }_{ii}W{h}_i+\sum _{j\in N(i)}{\alpha }_{ij}W{h}_j$$上述过程可以与下图对应：

我们可以多次计算自注意力分布并进行信息聚合，然后将不同自注意力空间下的结果融合即为多头注意力机制；

基于上述了解，现在再次回到GATConv的使用说明中：

Class GATConv(in_channels: int or Tuple(int, int), 
			out_channels: int, # 输出样本的特征维数
			heads: int = 1,  # 多头注意力的自注意力数量
			concat: bool = True, # 如果为False, 则对多头注意力输出特征进行平均而不是拼接
			negative_slope: float = 0.2, # LeakyReLU在负半轴上的负斜率
			dropout: float = 0.0, # 在训练期间采样邻居节点的丢弃率
			add_self_loops: bool = True, 
			bias: bool = True, 
			**kwargs)

关于第一个参数in_channels，与输入的图类型有关，对于常规的图，令in_channels=int；对于两个节点类型不同的图，使用in_channels=(int,int)，元组内的内容代表(源节点特征维数，目标节点特征维数)；

比如二分图，节点集可分割为两个互不相交的子集，并且图中每条边依附的两个节点都分别属于这两个互不相交的子集，两个子集内的节点不相邻：

这样的图，两个子集内的节点，其特征维数是不同的，所以需要利用元组输入到in_channels参数；

---

注意
在Pytorch-geometric的模型中，如果令输入特征维数in_channels=-1，表示延迟模型的参数初始化，模型会根据输入数据的形状自动初始化对应的参数形状

---

Using the Heterogenous Convolution Wrapper（示例）

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

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear

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):
            # 对不同类型的边使用不同的GCN
            conv = HeteroConv({
                ('paper', 'cites', 'paper'): GCNConv(-1, hidden_channels),
                ('author', 'writes', 'paper'): GATConv((-1, -1), hidden_channels),
                ('author', 'affiliated_with', 'institution'): SAGEConv((-1, -1), hidden_channels)
            }, 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 Heterogenous Operators（示例）

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

import os.path as osp

import torch
import torch.nn.functional as F
from torch_geometric.datasets import DBLP
from torch_geometric.nn import Linear, HGTConv

path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/DBLP')
dataset = DBLP(path)
data = dataset[0]
print(data.node_types) # ['author', 'paper', 'term', 'conference']

# 我们用同一个值初始化conference类型节点的特征
data['conference'].x = torch.ones(data['conference'].num_nodes, 1)

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=4, num_heads=2, num_layers=1)

异构图采样器Heterogeneous Graph Samplers

关于异构小批量加载：异构图可以分别通过 loader.DataLoader 和 loader.NeighborLoader 加载器转换为许多小型和单个巨型图的小批量。 这些加载器现在可以处理同构图和异构图：

# 处理多个图组成的集合
from torch_geometric.loader import DataLoader

loader = DataLoader(heterogeneous_graph_dataset, batch_size=32, shuffle=True)

# 处理单个巨型图
from torch_geometric.loader import NeighborLoader

loader = NeighborLoader(heterogeneous_graph, num_neighbors=[30, 30], batch_size=128,
                        input_nodes=('paper', data['paper'].train_mask), shuffle=True)

注意，在使用数据加载器 loader.NeighborLoader时，主要是对于巨型图的数据加载，我们最好符合以下版本（否则）：

torch                     1.9.0+cu111
torch-cluster             1.5.9
torch-geometric           2.0.0
torch-scatter             2.0.8
torch-sparse              0.6.12
torch-spline-conv         1.2.1

另外安装Python的扩展序列化/反序列化模块dill，dill将python用于序列化和反序列化python对象的pickle模块扩展到大多数内置python类型：

pip install dill

关于Win系统的dataloader，尽量让num_workers=0，下面是loader.NeighborLoader的用法；

loader.NeighborLoader是一个数据加载器，它执行“Inductive Representation Learning on Large Graphs”（大图上的归纳表示学习，回顾图神经网络第五课.可变图结构下的归纳式学习&图注意力）论文中介绍的邻居采样。 该加载器允许在无法进行全批量训练的大规模图上对 GNN 进行小批量训练；

更具体地说，num_neighbors 表示在每次迭代中为每个节点采样了多少邻居。 NeighborLoader 接收这个 num_neighbors 列表，为迭代第 i - 1 次中涉及到的每个节点再采样 num_neighbors[i]个对应的邻居节点；

常用参数如下：

CLASS NeighborLoader(
 # Data 或 HeteroData 图对象
 data: Union[torch_geometric.data.data.Data, torch_geometric.data.hetero_data.HeteroData], 

 # num_neighbors表示在每次迭代中为每个节点采样的邻居数, 在异构图中, 也可以使用一个字典来表示每个单独的边类型要采样的邻居数量
 num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]],
  
 # 对其邻居进行采样以创建小批量的节点的索引, 需要作为 torch.LongTensor 或 torch.BoolTensor 给出; 如果设置为 None, 则将考虑所有节点; 在异构图中, 需要作为保存节点类型和节点索引的元组传递 
 input_nodes: Union[torch.Tensor, None, str, Tuple[str, Optional[torch.Tensor]]] = None, 
 
 # 一个变换, 它接受一个采样的小批量并返回一个转换后的版本
 transform: Optional[Callable] = None, 

 # 额外的参数, 比如 batch_size, shuffle, num_workers
 **kwargs)

下面我们用NeighborLoader批量化加载一个巨型异构图：

import os.path as osp

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


path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/OGB')

transform = T.ToUndirected(merge=True)
dataset = OGB_MAG(path, preprocess='metapath2vec', transform=transform)
data = dataset[0]

print(data)

train_input_nodes = ('paper', data['paper'].train_mask)
val_input_nodes = ('paper', data['paper'].val_mask)

train_loader = NeighborLoader(data,
							  # 为每个节点采样 10 个邻居并进行 2 次迭代
                              num_neighbors=[10,10],
                              shuffle=True,
                              input_nodes=train_input_nodes,
                              batch_size=1024,
                              num_workers=0)

val_loader = NeighborLoader(data,
                            num_neighbors=[10,10],
                            input_nodes=val_input_nodes,
                            batch_size=1024,
                            num_workers=0)


print(train_loader) # NeighborLoader()

print(next(iter(train_loader)))
print(next(iter(train_loader)))

print(next(iter(train_loader))['paper'].batch_size) #1024

从CSV加载Graph

在此示例中，我们将学习如何加载一组 *.csv 文件作为输入并从中构建异构图，该图可用作异构图模型的输入。

我们使用 GroupLens 研究小组收集的 MovieLens 数据集。 这个数据集描述了来自 MovieLens 的 5 星评级和标记活动。 该数据集包含来自 600 多个用户的 9000 多部电影的大约 100k 评分。 我们将使用这个数据集生成两种分别保存电影和用户数据的节点类型，以及一种连接用户和电影的边类型，表示用户如何评价特定电影的关系。

首先，我们将数据集下载到任意文件夹（比如当前目录）：

from torch_geometric.data import download_url, extract_zip

url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
extract_zip(download_url(url, '.'), '.') # 下载并解压

movie_path = './ml-latest-small/movies.csv'
rating_path = './ml-latest-small/ratings.csv'

在我们创建异构图之前，让我们先看一下数据：

import pandas as pd

pd.read_csv(movie_path).head()

我们看到movies.csv 文件提供了三列：movieId 为每部电影分配一个唯一标识符，而title 和genres 列代表给定电影的标题和流派。 我们可以利用这两列来定义一个可以被机器学习模型感知的特征表达。

pd.read_csv(rating_path).head()

ratings.csv 数据连接用户（由 userId 给出）和电影（由 movieId 给出），并定义给定用户如何对特定电影进行评分（ rating ）。 为简单起见，我们不使用额外的时间戳信息（timestamp）。

为了以 PyG 数据格式表示此数据，我们首先定义了一个方法 load_node_csv()，该方法读取 *.csv 文件并返回形状为 [num_nodes, num_features] 的节点级特征表示 x：

import torch

def load_node_csv(path, index_col, encoders=None, **kwargs):
    df = pd.read_csv(path, index_col=index_col, **kwargs)
    """
    df=pd.read_csv(rating_path).head(5)
	df.index # RangeIndex(start=0, stop=5, step=1)
	df.index.unique() # Int64Index([0, 1, 2, 3, 4], dtype='int64')
    """
    mapping = {index: i for i, index in enumerate(df.index.unique())}

    x = None
    if encoders is not None:
    	# encoders.itmes() 每次迭代返回元组对象(key,element)
        xs = [encoder(df[col]) for col, encoder in encoders.items()]
        x = torch.cat(xs, dim=-1)

    return x, mapping

此处，load_node_csv() 从路径读取 *.csv 文件，并创建一个字典映射，将其索引列映射到 { 0, ..., num_rows - 1 } 范围内的连续值。这是必需的，因为我们希望我们的最终数据表示尽可能紧凑，例如，第一行中电影的表示应该可以通过 x[0] 访问。

我们进一步利用了编码器的概念，它定义了如何将特定列的值编码为数字特征表示。 例如，我们可以定义一个句子编码器，将原始列字符串编码为低维嵌入。 为此，我们利用了优秀的sentence-transformers，该库提供了大量最先进的预训练 NLP 嵌入模型，其安装如下：

pip install sentence-transformers

利用它进行编码：

from sentence_transformers import SentenceTransformer

class SequenceEncoder(object):
    def __init__(self, model_name='all-MiniLM-L6-v2', device=None):
        self.device = device
        self.model = SentenceTransformer(model_name, device=device)

    @torch.no_grad() # 禁用梯度, 前向计算节省显存
    def __call__(self, df):
        x = self.model.encode(df.values, show_progress_bar=True,
                              convert_to_tensor=True, device=self.device)
        return x.cpu()

SequenceEncoder 类加载由 model_name 给定的预训练 NLP 模型，并使用它将字符串列表编码为形状为 [num_strings, embedding_dim] 的 PyTorch 张量。 我们可以使用这个 SequenceEncoder 来编码movies.csv 文件的标题；

以类似的方式，我们可以创建另一个编码器，将电影类型（例如 Adventure|Children|Fantasy）转换为分类标签。 为此，我们首先需要找到数据中存在的所有现有流派，创建形状为 [num_movies, num_genres] 的特征表示 x，并为 x[i, j] 分配 1，以表示电影 i 中存在流派 j ：

class GenresEncoder(object):
    def __init__(self, sep='|'):
        self.sep = sep

    def __call__(self, df):
        genres = set(g for col in df.values for g in col.split(self.sep))
        mapping = {genre: i for i, genre in enumerate(genres)}

        x = torch.zeros(len(df), len(mapping)) # [num_movies, num_genres]
        for i, col in enumerate(df.values):
            for genre in col.split(self.sep):
                x[i, mapping[genre]] = 1
        return x

有了这个，我们可以通过以下方式获得电影的最终表示：

movie_x, movie_mapping = load_node_csv(
    movie_path, index_col='movieId', encoders={
        'title': SequenceEncoder(),
        'genres': GenresEncoder()
    })

同样，我们也可以利用 load_node_csv() 来获取从 userId 到连续值的用户映射。 但是，此数据集中的用户没有其他特征信息。 因此，我们没有定义任何编码器：

_, user_mapping = load_node_csv(rating_path, index_col='userId')

有了这个，我们准备初始化我们的 HeteroData 对象并将两种节点类型传递给它：

from torch_geometric.data import HeteroData

data = HeteroData()

data['user'].num_nodes = len(user_mapping)  # Users do not have any features.
data['movie'].x = movie_x

print(data)
"""
HeteroData(
  user={ num_nodes=610 },
  movie={ x[9742, 404] }
)
"""

由于用户没有任何节点级别的信息，我们仅定义其节点数。 因此，在训练异构图模型期间，我们可能需要通过 torch.nn.Embedding 以端到端的方式学习不同的用户嵌入。

接下来，我们看看将用户与他们的评级定义的电影联系起来。 为此，我们定义了一个方法 load_edge_csv()，该方法从 ratings.csv 返回形状 [2, num_ratings] 的最终 edge_index 表示，以及原始 *.csv 文件中存在的任何其他特征：

def load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping,
                  encoders=None, **kwargs):
    df = pd.read_csv(path, **kwargs)

    src = [src_mapping[index] for index in df[src_index_col]]
    dst = [dst_mapping[index] for index in df[dst_index_col]]
    edge_index = torch.tensor([src, dst])

    edge_attr = None
    if encoders is not None:
        edge_attrs = [encoder(df[col]) for col, encoder in encoders.items()]
        edge_attr = torch.cat(edge_attrs, dim=-1)

    return edge_index, edge_attr

与 load_node_csv() 类似，编码器用于返回额外的边级特征信息。 例如，为了从 ratings.csv 中的 rating 列加载 ratings，我们可以定义一个 IdentityEncoder 来简单地将浮点值列表转换为 PyTorch 张量：

class IdentityEncoder(object):
    def __init__(self, dtype=None):
        self.dtype = dtype

    def __call__(self, df):
        return torch.from_numpy(df.values).view(-1, 1).to(self.dtype)

有了这个，我们便可以准备好完成 HeteroData 对象：

edge_index, edge_label = load_edge_csv(
    rating_path,
    src_index_col='userId',
    src_mapping=user_mapping,
    dst_index_col='movieId',
    dst_mapping=movie_mapping,
    encoders={'rating': IdentityEncoder(dtype=torch.long)},
)

data['user', 'rates', 'movie'].edge_index = edge_index
data['user', 'rates', 'movie'].edge_label = edge_label

print(data)
"""
HeteroData(
  user={ num_nodes=610 },
  movie={ x=[9742, 404] },
  (user, rates, movie)={
    edge_index=[2, 100836],
    edge_label=[100836, 1]
  }
)
"""

这个 Hetero Data 对象是 PyG 中异构图的原生格式，可以用作异构图模型的输入