从CSV加载图表

在这个例子中，我们将展示如何加载一组*.csv文件作为输入，并从中构建一个异构图，该图可以用作异构图模型的输入。本教程也可作为可执行的示例脚本在examples/hetero目录中找到。

我们将使用由GroupLens研究小组收集的MovieLens数据集。这个玩具数据集描述了来自MovieLens的五星评分和标签活动。该数据集包含来自600多名用户对超过9千部电影的大约10万条评分。我们将使用这个数据集生成两种节点类型，分别保存电影和用户的数据，以及一种连接用户和电影的边类型，表示用户如何对特定电影进行评分的关系。

首先，我们将数据集下载到一个任意文件夹（在本例中为当前目录）：

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

print(pd.read_csv(movie_path).head())
print(pd.read_csv(rating_path).head())

Head of `movies.csv`
电影ID	标题	类型
1	玩具总动员 (1995)	冒险\|动画\|儿童\|喜剧\|奇幻
2	勇敢者的游戏 (1995)	冒险\|儿童\|奇幻
3	脾气更坏的老人们 (1995)	喜剧\|爱情
4	等待呼吸 (1995)	喜剧\|剧情\|爱情
5	新娘的父亲2 (1995)	喜剧

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

Head of `ratings.csv`
userId	电影ID	评分	时间戳
1	1	4.0	964982703
1	3	4.0	964981247
1	6	4.0	964982224
1	47	5.0	964983815
1	50	5.0	964982931

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)
    mapping = {index: i for i, index in enumerate(df.index.unique())}

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

    return x, mapping

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

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

pip install sentence-transformers

class SequenceEncoder:
    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 文件的 title。

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

class GenresEncoder:
    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))
        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

在这里，src_index_col 和 dst_index_col 分别定义了源节点和目标节点的索引列。我们进一步利用节点级别的映射 src_mapping 和 dst_mapping 来确保原始索引被映射到我们最终表示中的正确连续索引。对于文件中定义的每条边，它会在 src_mapping 和 dst_mapping 中查找前向索引，并适当地移动数据。

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

class IdentityEncoder:
    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]
  }
)

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

注意

点击这里查看最终的示例脚本。