官方文档部分

对于数据集，官方提供了两个抽象类：torch_geometric.data.Dataset和 torch_geometric.data.InMemoryDataset,torch_geometric.data.InMemoryDataset继承自torch_geometric.data.Dataset，如果整个数据集适合放入CPU内存中，则应该使用它。

每个数据集都会传递一个根文件夹，该文件夹指示了数据集应该存储的位置。我们将根文件夹分为两个文件夹：raw_dir，用于下载数据集；processed_dir，用于保存处理过的数据集。

此外，每个数据集都可以传递一个transoform(变换？处理)函数、一个pre_transform(预处理函数)和一个pre_filter(预过滤)函数，这些默认情况下都是None。

1. transform函数在访问之前动态地转换数据对象（因此它最适合用于数据增强）。
2. pre_transform函数在将数据对象保存到磁盘之前应用转换（因此它最适合用于需要只进行一次的重度的预计算）。
3. pre_filter函数可以在保存之前手动过滤掉数据对象。

创建一个存入内存的数据集

为了创建一个torch_geometric.data.InMemoryDataset，你需要实现四个基本方法：

1. InMemoryDataset.raw_file_names():需要找到的raw_dir中的文件列表，以便跳过下载。
2. InMemoryDataset.processed_file_names():需要找到的processed_dir中的文件列表，以便跳过处理。
3. InMemoryDataset.download():下载raw数据进入raw_dir目录。
4. InMemoryDataset.process():处理raw数据并且保存进入processed_dir目录。

可以在torch_geometric.data这个类里面学习相关例子。

真正的魔法发生在process()函数的体内。在这里，我们需要读取并创建一个Data对象的列表，并将其保存到processed_dir中。由于保存一个巨大的Python列表相当慢，我们在保存之前通过torch_geometric.data.InMemoryDataset.collate()将列表合并成一个巨大的Data对象。合并后的数据对象将所有示例连接成一个大的数据对象，并且，还会返回一个切片字典，以便从这个对象重建单个示例。最后，我们需要在构造函数(constructor)中将这两个对象加载到属性self.data和self.slices中。

self.data和self.slices通过torch_geometric.data.InMemoryDataset.load()隐式的加载。
注意：在 PyG ​​版本 < 2.4​​ 时，需要使用 torch.load(self.processed_paths[0]) 然后手动拆解成 data 和 slices，但在 ​​PyG >= 2.4​​ 后，直接使用 self.load(path) 即可，内部已经帮你处理好了。

在 PyTorch Geometric（PyG）的 InMemoryDataset 类中，self.processed_paths 是一个属性，它返回一个 ​​列表​​，包含所有处理后的数据文件（即经过 process() 方法处理并保存的数据）的路径。

官方给了一个例子

# 加载torch，加载pyg的data.InMemoryDataset模块
import torch
from torch_geometric.data import InMemoryDataset, download_url

# 创建一个自定义数据集，继承InMemoryDataset类，并且实现上述四种函数，若不实现将会报错。
class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super().__init__(root, transform, pre_transform, pre_filter)
        self.load(self.processed_paths[0])
        # For PyG<2.4:
        # self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['some_file_1', 'some_file_2', ...]

    @property
    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        # Download to `self.raw_dir`.
        download_url(url, self.raw_dir)
        ...

    def process(self):
        # Read data into huge `Data` list.
        data_list = [...]

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        self.save(data_list, self.processed_paths[0])
        # For PyG<2.4:
        # torch.save(self.collate(data_list), self.processed_paths[0])

创建一个“更大”的数据集

当数据集在内存里存不下的情况下，我们将使用torch_geometric.data.Dataset,它需要额外实现如下两个方法：

1. Dataset.len():返回数据集中示例数量。
2. Dataset.get():实现加载单个graph的逻辑。

torch_geometric.data.Dataset.__getitem__()函数从torch_geometric.data.Dataset.get() 函数获取数据对象，并根据transform进行可选的转换。

import os.path as osp

import torch
from torch_geometric.data import Dataset, download_url


class MyOwnDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super().__init__(root, transform, pre_transform, pre_filter)

    @property
    def raw_file_names(self):
        return ['some_file_1', 'some_file_2', ...]

    @property
    def processed_file_names(self):
        return ['data_1.pt', 'data_2.pt', ...]

    def download(self):
        # Download to `self.raw_dir`.
        path = download_url(url, self.raw_dir)
        ...

    def process(self):
        idx = 0
        for raw_path in self.raw_paths:
            # Read data from `raw_path`.
            data = Data(...)

            if self.pre_filter is not None and not self.pre_filter(data):
                continue

            if self.pre_transform is not None:
                data = self.pre_transform(data)

            torch.save(data, osp.join(self.processed_dir, f'data_{idx}.pt'))
            idx += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, f'data_{idx}.pt'))
        return data

其他注意事项

1. 跳过download()和/或process():
   通过不覆盖上述函数即可：

class MyOwnDataset(Dataset):
    def __init__(self, transform=None, pre_transform=None):
        super().__init__(None, transform, pre_transform)

2. 是否可以不创建databasets？

可以，直接使用Data和DataLoader对象：

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

data_list = [Data(...), ..., Data(...)]
loader = DataLoader(data_list, batch_size=32)

从CSV文件中加载图像

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

我们将使用GroupLens研究小组收集的电影透镜数据集。这个玩具数据集描述了电影透镜的5星评级和标记活动。该数据集包含大约100k个评分，涉及超过9k部电影，来自超过600名用户。我们将使用这个数据集生成两种节点类型，分别用于存储电影和用户的数据，以及一种边类型，连接用户和电影，代表用户对特定电影的评级关系。

我们首先将数据集下载到任意文件夹下，例如当前目录:

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())

内容如下：

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

ratings.csv数据连接了用户（由userId给出）和电影（由movieId给出），并定义了给定用户对特定电影的评分（rating）。由于简单起见，我们没有使用额外的timestamp信息。

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

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() 从路径读取 *.csv文件，并创建一个字典映射，将其索引列映射到范围 { 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类加载一个预训练的NLP模型，如model_name所示，并使用它来将字符串列表编码成形状为[num_strings, embedding_dim]的PyTorch张量。我们可以使用这个SequenceEncoder来编码movies.csv文件的标题。

以类似的方式，我们可以创建另一个编码器，将电影的类型（例如，冒险|儿童|幻想）转换为分类标签。为此，我们首先要找到所有的电影类型，创建一个形状为[num_movies, num_genres]的特征表示x，并在电影i中存在类型j的情况下将1分配给x[i, j]：

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返回最终的edge_index表示，形状为[2, num_ratings]。

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中的评分列加载评分，我们可以定义一个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]
  }
)

数据集划分

数据集划分是图机器学习中的一个关键步骤，我们将我们的数据集划分为训练集、验证集和测试集。这确保了我们的模型得到适当的评估，防止过拟合，并实现泛化。在本教程中，我们将探讨数据集划分的基础知识，重点关注三个基本任务：节点预测、链接预测和图预测。我们将介绍常用的技术，包括RandomNodeSplit和RandomLinkSplit变换。此外，我们还将涵盖如何创建自定义的数据集划分，而不仅仅是随机的划分。

Node划分

在本节中，我们将学习如何使用PyTorch Geometric库中的RandomNodeSplit函数将节点随机划分为训练集、验证集和测试集。在examples/cora.py中提供了一个完整的工作示例，该示例使用了Planetoid数据集。

RandomNodeSplit被初始化为用于分割PyG Data和HeteroData对象的节点。

• split 定义dataset's分割类型
• num_split定义要添加的split数目
• num_train_per_class 定义每个类别的训练节点数。
• num_val验证节点数
• num_test测试节点数
• key 真实标签名

import torch
from torch_geometric.data import Data
from torch_geometric.transforms import RandomNodeSplit

x = torch.randn(8, 32)  # Node features of shape [num_nodes, num_features]
y = torch.randint(0, 4, (8, ))  # Node labels of shape [num_nodes]
edge_index = torch.tensor([
    [2, 3, 3, 4, 5, 6, 7],
    [0, 0, 1, 1, 2, 3, 4]],
)

#   0  1
#  / \/ \
# 2  3  4
# |  |  |
# 5  6  7

data = Data(x=x, y=y, edge_index=edge_index)
node_transform = RandomNodeSplit(num_val=2, num_test=3)
node_splits = node_transform(data)

在这里，我们初始化一个RandomNodeSplit转换来按节点分割图数据。转换后，train_mask、valid_mask和test_mask将被附加到图数据上。

node_splits.train_mask
>>> tensor([ True, False, False, False, True, True, False, False])
node_splits.val_mask
>>> tensor([False, False, False, False, False, False, True, True])
node_splits.test_mask
>>> tensor([False, True, True, True, False, False, False, False])

在这个例子中，有8个节点，我们想要采样2个节点用于验证，3个节点用于测试，其余的用于训练。最终，我们得到了节点0、4、5作为训练集，节点6、7作为验证集，以及节点1、2、3作为测试集。