基于Pytorch的PyTorch Geometric（PYG）库构造个人数据集

一、目的

PyTorch Geometric库中虽然已经包含自带的数据集如 Cora 等，但有时我们也需要用户个人数据创建自己的数据集进行一些数据研究。当然博主也建议大家若是第一次使用PyTorch Geometric库可以先使用其自带的数据集进行理解，再创建自己的数据集做到灵活运用。

二、前期准备和参考资料

三、PyTorch Geometric库构造个人数据集三种实现方式

PyG为数据集提供了三个抽象类：Data、InMemoryDataset 和 Dataset。

  • torch_geometric.data.Data：用户只需要在程序运行中动态创建合成数据且不需要将数据以pt文件形式保存到磁盘时，直接可以在读取所有外部文件数据集后，只需传递一个包含torch_geometry.data.Data的常规python列表数据对象，并将其传递给torch_geometry.Data.dataloader即可进行后续数据处理操作，一般适合小规模数据集，用户个人数据集可以选这个模式处理即可，且博主的 这篇《基于Pytorch的图卷积网络GCN实例应用及详解》文章 构造个人数据集使用的便是此类方法。
  • torch_geometric.data.InMemoryDataset：是继承torch_geometry.data.Dataset，并且如果整个数据集都能装入内存，则可以使用它，一般适合小规模数据集，用户个人数据集可以选这个模式处理即可。
  • torch_geometry.data.Dataset：需要每读取一个Batch数据会重新从硬盘中读取这些数据再放入程序运行的缓存中，一般比较适合大规模的数据集（此篇文章不详细介绍）。

四、torch_geometric.data.InMemoryDataset基本方法

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

  • torch_geometric.data.InMemoryDataset.raw_file_names()：首先寻找raw_dir其中的文件列表（见下图raw），文件为空时，才需要下载download()数据集，若想跳过此方法则返回一个字符串或者字节即可。
  • torch_geometric.data.InMemoryDataset.processed_file_names()：首先寻找processed_dir其中的文件列表（见下图processed），文件为空时，才需要进行数据集process()处理。
  • torch_geometric.data.InMemoryDataset.download(): 将原始数据下载到raw_dir（见下图raw），当用户自己创建数据集时也可使用pass方法跳过此过程。
  • torch_geometric.data.InMemoryDataset.process()：处理原始数据集并将其保存到processed_dir（见下图processed）。在数据处理函数中，具体流程是：用户需要读取并创建一个Data对象列表并将其保存到processed_dir.。因为保存一个巨大的 python 列表速度非常慢，因此在保存前可以通过 torch_geometric.data.InMemoryDataset.collate() 方法将列表整理成一个巨大的Data对象，再返回一个Data对象切片（slices）字典和一个Data对象重建的单个示例，最后我们需要在构造函数中将这两个对象加载到属性 self.data 和 self.slices 。

五、torch_geometric.data.InMemoryDataset实现流程代码及结果

InMemoryDataset实现流程图如下：

完整代码如下：

import torch
from torch_geometric.data import InMemoryDataset, download_url
from torch_geometric.data import Data
from torch_geometric.data import DataLoader
import warnings
warnings.filterwarnings("ignore", category=Warning)
from torch_geometric.datasets import TUDataset
# 这里给出大家注释方便理解
# 程序只要第一次运行后，processed文件生成后就不会执行proces函数，而且只要不重写download()和process()方法，也会直接跳过下载和处理。
class MyOwnDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])
        print(self.data) # 输出torch.load加载的数据集data
        # print(root) # MYdata
        # print(self.data) # Data(x=[3, 1], edge_index=[2, 4], y=[3])
        # print(self.slices) # defaultdict(<class 'dict'>, {'x': tensor([0, 3, 6]), 'edge_index': tensor([ 0,  4, 10]), 'y': tensor([0, 3, 6])})
        # print(self.processed_paths[0]) # MYdata\processed\datas.pt

    # 返回数据集源文件名，告诉原始的数据集存放在哪个文件夹下面，如果数据集已经存放进去了，那么就会直接从raw文件夹中读取。
    @property
    def raw_file_names(self):
        # pass # 不能使用pass，会报join() argument must be str or bytes, not 'NoneType'错误
        return []

    # 首先寻找processed_paths[0]路径下的文件名也就是之前process方法保存的文件名
    @property
    def processed_file_names(self):
        return ['datas.pt']

    # 用于从网上下载数据集，下载原始数据到指定的文件夹下，自己的数据集可以跳过
    def download(self):
        pass

    # 生成数据集所用的方法，程序第一次运行才执行并生成processed文件夹的处理过后数据的文件，否则必须删除已经生成的processed文件夹中的所有文件才会重新执行此函数
    def process(self):
        # Read data into huge `Data` list.
        # Read data into huge `Data` list.
        # 这里用于构建data
        edge_index1 = torch.tensor([[0, 1, 1, 2],
                                   [1, 0, 2, 1]], dtype=torch.long)

        edge_index2 = torch.tensor([[0, 1, 1, 2 ,0 ,1],
                                    [1, 0, 2, 1 ,0 ,1]], dtype=torch.long)

        # 节点及每个节点的特征：从0号节点开始
        X = torch.tensor([[-1], [0], [1]], dtype=torch.float)
        # 每个节点的标签：从0号节点开始-两类0，1
        Y = torch.tensor([0, 1, 0], dtype=torch.float)

        # 创建data数据
        data1 = Data(x=X, edge_index=edge_index1, y=Y)
        data2 = Data(x=X, edge_index=edge_index2, y=Y)

        # 将data放入datalist
        data_list = [data1,data2]
        # data_list = data_list.append(data)

        if self.pre_filter is not None: # pre_filter函数可以在保存之前手动过滤掉数据对象。用例可能涉及数据对象属于特定类的限制。默认None
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None: # pre_transform函数在将数据对象保存到磁盘之前应用转换(因此它最好用于只需执行一次的大量预计算)，默认None
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list) # 直接保存list可能很慢，所以使用collate函数转换成大的torch_geometric.data.Data对象
        # print(data)
        torch.save((data, slices), self.processed_paths[0])

# 数据集对象操作
b = MyOwnDataset("E:\pycharmproject\MYdata") # 创建数据集对象
data_loader = DataLoader(b, batch_size=1, shuffle=False) # 加载数据进行处理，每批次数据的数量为1
for data in data_loader:
    print(data) # 按批次输出数据

程序运行结果如下：

程序代码注意三事项：

  • raw_file_names()：此方法必须返回一个字符串或者字节，不能直接使用pass不传递任何值，否则会报 join() argument must be str or bytes, not ‘NoneType’ 错误。
  • download()：此方法在用户自己创建数据集的情况下，一定要直接pass跳过，因为不需要下载。
  • process()：此方法是程序只要第一次运行后，processed文件（processed_dir）自动生成后，就不会再执行原程序的process方法，之后每次只会提取processed_dir文件的数据集，因此只要不重写process()方法和将原来的processed_dir删除掉，程序第二次运行开始会直接跳过处理步骤。

六、torch_geometry.data.Dataset（了解即可）

相较于torch_geometric.data.InMemoryDataset基础上额外增加两个方法：

torch_geometric.data.Dataset.len()： 返回数据集中示例的数目。
torch_geometric.data.Dataset.get():：实现加载单个图的逻辑。

# 返回数据集中示例的数目
def len(self):
    return len(self.processed_file_names)
# 实现加载单个图的逻辑
def get(self, idx):
    data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
    return data