（DataWhale）图神经网络Task07：按需加载样本到内存的数据集类

背景：当数据集规模超级大，很难在内存中完全存下所有数据，因此要按需将数据按需加载到内存。

简单数据导入

依赖torch_geometric.data.Data, DataLoader的数据导入

from torch_geometric.data import Data, DataLoader

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

依赖torch_geometric.data.Data, Batch的批次导入

from torch_geometric.data import Data, Batch

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

自定义数据导入

继承Dataset基类的自定义数据集类

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):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['raw_file_1', 'raw_file_2', ...]

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

    def download(self):
        path = download_url(url, self.raw_dir)

    def process(self):
        ...

    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

其中，len()：返回数据集中的样本的数量；get()：实现加载单个图的操作，在内部，__getitem__()返回通过调用get()来获取Data对象，并根据transform参数对它们进行选择性转换。

图样本封装成批（BATCHING）

PyTorch Geometric中采用的将多个图封装成批的方式是，将小图作为连通组件（connected component）的形式合并，构建一个大图。于是小图的邻接矩阵存储在大图邻接矩阵的对角线上。大图的邻接矩阵、属性矩阵、预测目标矩阵分别为：
KaTeX parse error: No such environment: split at position 8: \begin{̲s̲p̲l̲i̲t̲}̲\mathbf{A} = \b…
此方法有以下关键的优势：

• 依靠消息传递方案的GNN运算不需要被修改，因为消息仍然不能在属于不同图的两个节点之间交换。
• 没有额外的计算或内存的开销。例如，这个批处理程序的工作完全不需要对节点或边缘特征进行任何填充。请注意，邻接矩阵没有额外的内存开销，因为它们是以稀疏的方式保存的，只保留非零项，即边。

通过torch_geometric.data.DataLoader类，多个小图被封装成一个大图。torch_geometric.data.DataLoader是PyTorch的DataLoader的子类，覆盖了collate()函数，该函数定义了一列表的样本是如何封装成批的。

小图的属性增值与拼接

将小图存储到大图中时需要对小图的属性做一些修改，一个最显著的例子就是要对节点序号增值。在最一般的形式中，PyTorch Geometric的DataLoader类会自动对edge_index张量增值，增加的值为当前被处理图的前面的图的累积节点数量。

默认操作是：

• 对第$k$个图的edge_index张量“前面$k-1$个图的累积节点数量$n$”的增量；
• 增值后，在第二维中连接所有图的edge_index张量（[2, num_edges]）。

通过覆盖torch_geometric.data.__inc__()和torch_geometric.data.__cat_dim__()函数可以根据实际需求更改默认操作。其中，__inc__()定义了两个连续的图的属性之间的增量大小，而__cat_dim__()定义了同一属性的图形张量应该在哪个维度上被连接起来。

图的匹配（Pairs of Graphs）

背景：在一个Data对象中存储多个图，例如在图匹配等应用中，需要将若干图正确封装成批。

以下PairData类将一个源图$G_s$和一个目标图$G_t$存储在Data类中：

class PairData(Data):
    def __init__(self, edge_index_s, x_s, edge_index_t, x_t):
        super(PairData, self).__init__()
        self.edge_index_s = edge_index_s
        self.x_s = x_s
        self.edge_index_t = edge_index_t
        self.x_t = x_t

    def __inc__(self, key, value):
        if key == 'edge_index_s':  # 根据源图G_s的节点数做增值
            return self.x_s.size(0)
        if key == 'edge_index_t':  # 根据目标图G_t的节点数做增值
            return self.x_t.size(0)
        else:
            return super().__inc__(key, value)

上述PairData类批处理行为的测试如下：

edge_index_s = torch.tensor([
    [0, 0, 0, 0],
    [1, 2, 3, 4],
])
x_s = torch.randn(5, 16)  # 5 nodes.
edge_index_t = torch.tensor([
    [0, 0, 0],
    [1, 2, 3],
])
x_t = torch.randn(4, 16)  # 4 nodes.

data = PairData(edge_index_s, x_s, edge_index_t, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)
# Batch(edge_index_s=[2, 8], x_s=[10, 16], edge_index_t=[2, 6], x_t=[8, 16])

print(batch.edge_index_s)
# tensor([[0, 0, 0, 0, 5, 5, 5, 5], [1, 2, 3, 4, 6, 7, 8, 9]])

print(batch.edge_index_t)
# tensor([[0, 0, 0, 4, 4, 4], [1, 2, 3, 5, 6, 7]])

由于PyTorch Geometric无法识别PairData对象中实际的图，所以batch属性（将大图每个节点映射到其各自对应的小图）没有正确工作。此时就需要DataLoader的follow_batch参数发挥作用，通过该参数可以指定要为哪些属性维护批信息。

loader = DataLoader(data_list, batch_size=2, follow_batch=['x_s', 'x_t'])
batch = next(iter(loader))

print(batch)
# Batch(edge_index_s=[2, 8], x_s=[10, 16], x_s_batch=[10],
#       edge_index_t=[2, 6], x_t=[8, 16], x_t_batch=[8])
print(batch.x_s_batch)
# tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])

print(batch.x_t_batch)
# tensor([0, 0, 0, 0, 1, 1, 1, 1])

可以看到，follow_batch=['x_s', 'x_t']成功地为节点特征x_s'和x_t’分别创建了名为x_s_batch和x_t_batch的赋值向量。这些信息可以用来在一个单一的Batch对象中对多个图进行聚合操作，例如全局池化。

二部图（Bipartite Graphs）

二部图的邻接矩阵定义两种类型的节点之间的连接关系。一般来说，不同类型的节点数量不需要一致，于是二部图的邻接矩阵 A ∈ { 0 , 1 } N × M A \in \{0,1\}^{N \times M} A∈{0,1}N×M可能为平方矩阵，即可能有$N\ne M$。在二部图的封装成批过程中，edge_index 中边的源节点与目标节点做的增值操作应是不同的。

class BipartiteData(Data):
    def __init__(self, edge_index, x_s, x_t):
        super(BipartiteData, self).__init__()
        self.edge_index = edge_index
        self.x_s = x_s
        self.x_t = x_t
        
    def __inc__(self, key, value):
        if key == 'edge_index':  # 在edge_index中独立地为边的源节点和目标节点做增值操作
        	return torch.tensor([[self.x_s.size(0)], [self.x_t.size(0)]])
    	else:
        	return super().__inc__(key, value)     

上述BipartiteData类批处理行为的测试如下：

edge_index = torch.tensor([
    [0, 0, 1, 1],
    [0, 1, 1, 2],
])
x_s = torch.randn(2, 16)  # 2 nodes.
x_t = torch.randn(3, 16)  # 3 nodes.

data = BipartiteData(edge_index, x_s, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)
# Batch(edge_index=[2, 8], x_s=[4, 16], x_t=[6, 16])

print(batch.edge_index)
# tensor([[0, 0, 1, 1, 2, 2, 3, 3], [0, 1, 1, 2, 3, 4, 4, 5]])

在新的维度上做拼接

Data对象的属性需要在一个新的维度上做拼接（经典的封装成批），例如图级别属性或预测目标。具体来说，形状为[num_features]的属性列表应该被返回为[num_examples, num_features]，而不是[num_examples * num_features]。PyTorch Geometric通过在__cat_dim__()中返回一个None的连接维度来实现这一点。

class MyData(Data):     def __cat_dim__(self, key, item):         if key == 'foo':             return None         else:             return super().__cat_dim__(key, item)edge_index = torch.tensor([   [0, 1, 1, 2],   [1, 0, 2, 1],])foo = torch.randn(16)data = MyData(edge_index=edge_index, foo=foo)data_list = [data, data]loader = DataLoader(data_list, batch_size=2)batch = next(iter(loader))print(batch)# Batch(edge_index=[2, 8], foo=[2, 16])# batch.foo中，2为批维度，16为特征维度

超大规模数据集类实践

PCQM4M-LSC是一个分子图的量子特性回归数据集，包含了3,803,453个图。

import osimport os.path as ospimport pandas as pdimport torchfrom ogb.utils import smiles2graphfrom ogb.utils.torch_util import replace_numpy_with_torchtensorfrom ogb.utils.url import download_url, extract_zipfrom rdkit import RDLoggerfrom torch_geometric.data import Data, Datasetimport shutilRDLogger.DisableLog('rdApp.*')class MyPCQM4MDataset(Dataset):    def __init__(self, root):        self.url = 'https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m_kddcup2021.zip'        super(MyPCQM4MDataset, self).__init__(root)        filepath = osp.join(root, 'raw/data.csv.gz')        data_df = pd.read_csv(filepath)        self.smiles_list = data_df['smiles']        self.homolumogap_list = data_df['homolumogap']    @property    def raw_file_names(self):        return 'data.csv.gz'    def download(self):        path = download_url(self.url, self.root)        extract_zip(path, self.root)        os.unlink(path)        shutil.move(osp.join(self.root, 'pcqm4m_kddcup2021/raw/data.csv.gz'), osp.join(self.root, 'raw/data.csv.gz'))    def len(self):        return len(self.smiles_list)    def get(self, idx):        smiles, homolumogap = self.smiles_list[idx], self.homolumogap_list[idx]        graph = smiles2graph(smiles)        assert(len(graph['edge_feat']) == graph['edge_index'].shape[1])        assert(len(graph['node_feat']) == graph['num_nodes'])        x = torch.from_numpy(graph['node_feat']).to(torch.int64)        edge_index = torch.from_numpy(graph['edge_index']).to(torch.int64)        edge_attr = torch.from_numpy(graph['edge_feat']).to(torch.int64)        y = torch.Tensor([homolumogap])        num_nodes = int(graph['num_nodes'])        data = Data(x, edge_index, edge_attr, y, num_nodes=num_nodes)        return data    # 获取数据集划分    def get_idx_split(self):        split_dict = replace_numpy_with_torchtensor(torch.load(osp.join(self.root, 'pcqm4m_kddcup2021/split_dict.pt')))        return split_dictif __name__ == "__main__":    dataset = MyPCQM4MDataset('dataset2')    from torch_geometric.data import DataLoader    from tqdm import tqdm    dataloader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=4)    for batch in tqdm(dataloader):        pass

• 在生成一个该数据集类的对象时，程序首先会检查指定的文件夹下是否存在data.csv.gz文件，如果不在，则会执行download方法，这一过程是在运行super类的__init__方法中发生的；
• 然后程序继续执行__init__方法的剩余部分，读取data.csv.gz文件，获取存储图信息的smiles格式的字符串，以及回归预测的目标homolumogap。

参考

1. DataWhale GNN组队学习