参考链接:https://github.com/datawhalechina/team-learning-nlp/blob/master/GNN
第一部分 超大规模数据集类的创建
当数据集规模超级大时,很难有足够大的内存完全存下所有数据。因此需要一个按需加载样本到内存的数据集类。
一、Dataset基类
1.1Dataset基类简介
在PyG中,通过继承torch_geometric.data.Dataset
基类来自定义一个按需加载样本到内存的数据集类。
继承此基类比继承torch_geometric.data.InMemoryDataset基类要多实现以下方法:
- len():返回数据集中的样本的数量。
- get():实现加载单个图的操作。注意:在内部,getitem()返回通过调用get()来获取Data对象,并根据transform参数对它们进行选择性转换。
1.2继承torch_geometric.data.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 ['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):
i = 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, 'data_{}.pt'.format(i)))
i += 1
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
1.3其他注意事项
1.download/process步骤可以跳过
- 对于无需下载数据集原文件的情况,不重写(override)download方法即可跳过下载。
- 对于无需对数据集做预处理的情况,不重写process方法即可跳过预处理。
2.有些Dataset类无需定义
如下,可以不用定义一个Dataset类,而直接生成一个Dataloader对象,直接用于训练。
from torch_geometric.data import Data, DataLoader
data_list = [Data(...), ..., Data(...)]
loader = DataLoader(data_list, batch_size=32)
二、图样本封装成批(BATCHING)与DataLoader类
2.1合并小图组成大图
PyTorch Geometric中采用的是将多个图封装成批的方式,将小图作为连通组件(connected component)的形式合并,构建一个大图。于是小图的邻接矩阵存储在大图邻接矩阵的对角线上。
此方法有以下关键的优势:
- 依靠消息传递方案的GNN运算不需要被修改。
- 没有额外的计算或内存的开销。
通过torch_geometric.data.DataLoader
类,多个小图被封装成一个大图。torch_geometric.data.DataLoader
是PyTorch的DataLoader
的子类,覆盖了collate()
函数,该函数定义了一列表的样本是如何封装成批的。因此,所有可以传递给PyTorch DataLoader
的参数也可以传递给PyTorch Geometric的 DataLoader
。
2.2小图的属性增值与拼接
将小图存储到大图中时需要对小图的属性做一些修改,一个最显著的例子就是要对节点序号增值。在最一般的形式中,PyTorch Geometric的DataLoader类会自动对edge_index张量增值,增加的值为当前被处理图的前面的图的累积节点数量。增值后,对所有图的edge_index张量(其形状为[2, num_edges])在第二维中连接起来。
2.2.1图的匹配(Pairs of Graphs)
不同类型的节点数量不一致,edge_index边的源节点与目标节点进行增值操作不同。
2.2.2二部图(Bipartite Graphs)
二部图是图论中的一种特殊模型。设G=(V,E)是一个无向图,如果顶点V可分割为两个互不相交的子集(A,B),并且图中的每条边(i,j)所关联的两个顶点i和j分别属于这两个不同的顶点集(i in A,j in B),则称图G为一个二部图。它的邻接矩阵定义两种类型的节点之间的连接关系。一般来说,不同类型的节点数量不需要一致,于是二部图的邻接矩阵可能为平方矩阵,即可能有。
2.2.3在新的维度上做拼接
有时,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
现在由两个维度来表示,一个批维度,一个特征维度。
三、创建超大规模数据集类实践
PCQM4M-LSC是一个分子图的量子特性回归数据集,它包含了3,803,453个图。
定义的数据集类如下:
import os
import os.path as osp
import pandas as pd
import torch
from ogb.utils import smiles2graph
from ogb.utils.torch_util import replace_numpy_with_torchtensor
from ogb.utils.url import download_url, extract_zip
from rdkit import RDLogger
from torch_geometric.data import Data, Dataset
import shutil
RDLogger.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_dict
if __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
。由smiles
格式的字符串转成图的过程在get()
方法中实现,这样在生成一个DataLoader
变量时,通过指定num_workers
可以实现并行执行生成多个图。
第二部分 图预测任务实践
1.通过试验寻找最佳超参数
通过运行以下的命令即可运行一次试验:
#!/bin/sh
python main.py --task_name GINGraphPooling\ # 为当前试验取名
--device 0\
--num_layers 5\ # 使用GINConv层数
--graph_pooling sum\ # 图读出方法
--emb_dim 256\ # 节点嵌入维度
--drop_ratio 0.\
--save_test\ # 是否对测试集做预测并保留预测结果
--batch_size 512\
--epochs 100\
--weight_decay 0.00001\
--early_stop 10\ # 当有`early_stop`个epoches验证集结果没有提升,则停止训练
--num_workers 4\
--dataset_root dataset # 存放数据集的根目录
试验运行开始后,程序会在saves
目录下创建一个task_name
参数指定名称的文件夹用于记录试验过程,当saves
目录下已经有一个同名的文件夹时,程序会在task_name
参数末尾增加一个后缀作为文件夹名称。试验运行过程中,所有的print
输出都会写入到试验文件夹下的output
文件,tensorboard.SummaryWriter
记录的信息也存储在试验文件夹下的文件中。
修改上方的命令再执行,即可试验不同的超参数,所有试验的过程与结果信息都存储于saves
文件夹下。启动TensorBoard
会话,选择saves
文件夹,即可查看所有试验的过程与结果信息。
2.总结
在此图预测任务实践中:
- 此次将前面所学的基于GIN的图表示学习神经网络和超大规模数据集类的创建方法付诸于实际应用;
- 构建了一种很方便的设置不同参数进行试验的方法,不同试验的过程与结果信息通过简单的操作即可进行比较分析。