第七天深入学习DGL框架：官方文档指导下的数据集下载与预处理技巧

参考链接

1. https://docs.dgl.ai/guide/data.html#guide-data-pipeline
2. https://docs.dgl.ai/en/0.5.x/_modules/dgl/data/qm7b.html#QM7bDataset

DGLDataset

DGL在 dgl.data 里实现了很多常用的图数据集。它们遵循了由 dgl.data.DGLDataset 类定义的标准的数据处理管道。DGL推荐用户将图数据处理为 dgl.data.DGLDataset 的子类。该类为导入、处理和保存图数据提供了简单而清晰的解决方案。

DGLDataset的执行流程&＃xff1a;

1. 通过调用“has_cache()”判断磁盘上是否有已经处理好的数据集缓存。如果有&＃xff0c;则跳转到第5步&＃xff0c;直接加载数据集&＃xff1b;

2. 调用“download()”下载数据&＃xff1b;

3. 调用“process()”处理数据&＃xff1b;

4. 调用“save()”保存处理好的数据到磁盘&＃xff0c;跳转到第6步&＃xff1b;

5. 调用“load()”从磁盘加载数据集&＃xff1b;

6. 完成。

下面给出了一个继承自DGLDataset类的例子。子类中必须实现process()&＃xff0c; getitem(idx) 和 len()。同时官方建议也实现save()和load()&＃xff0c;避免对大型数据集的重复处理。

from dgl.data import DGLDatasetclass MyDataset(DGLDataset):""" 用于在DGL中自定义图数据集的模板&＃xff1a;Parameters----------url : str下载原始数据集的url。raw_dir : str指定下载数据的存储目录或已下载数据的存储目录。默认: ~/.dgl/save_dir : str处理完成的数据集的保存目录。默认&＃xff1a;raw_dir指定的值force_reload : bool是否重新导入数据集。默认&＃xff1a;Falseverbose : bool是否打印进度信息。"""def __init__(self,url&＃61;None,raw_dir&＃61;None,save_dir&＃61;None,force_reload&＃61;False,verbose&＃61;False):super(MyDataset, self).__init__(name&＃61;&＃39;dataset_name&＃39;,url&＃61;url,raw_dir&＃61;raw_dir,save_dir&＃61;save_dir,force_reload&＃61;force_reload,verbose&＃61;verbose)def download(self):# 将原始数据下载到本地磁盘passdef process(self):# 将原始数据处理为图、标签和数据集划分的掩码passdef __getitem__(self, idx):# 通过idx得到与之对应的一个样本passdef __len__(self):# 数据样本的数量passdef save(self):# 将处理后的数据保存至 &＃96;self.save_path&＃96;passdef load(self):# 从 &＃96;self.save_path&＃96; 导入处理后的数据passdef has_cache(self):# 检查在 &＃96;self.save_path&＃96; 中是否存有处理后的数据pass


下载原始数据

这一段就是给实现“download()”举了两个例子。
从“self.url”链接下载到“self.raw_dir”目录下&＃xff0c;保存为“self.name&＃43;格式后缀”&＃xff1a;

import os
from dgl.data.utils import downloaddef download(self):# 存储文件的路径file_path &＃61; os.path.join(self.raw_dir, self.name &＃43; &＃39;.mat&＃39;)# 下载文件download(self.url, path&＃61;file_path)


如果数据集是一个zip文件&＃xff0c;可以直接继承 dgl.data.DGLBuiltinDataset 类&＃xff0c;其支持解压缩zip文件。

如果文件是.gz、.tar、.tar.gz或.tgz文件&＃xff0c;下载后需要用 extract_archive() 函数进行解压缩&＃xff1a;

from dgl.data.utils import download, check_sha1def download(self):# 存储文件的路径&＃xff0c;请确保使用与原始文件名相同的后缀gz_file_path &＃61; os.path.join(self.raw_dir, self.name &＃43; &＃39;.csv.gz&＃39;)# 下载文件download(self.url, path&＃61;gz_file_path)# 检查 SHA-1if not check_sha1(gz_file_path, self._sha1_str):raise UserWarning(&＃39;File {} is downloaded but the content hash does not match.&＃39;&＃39;The repo may be outdated or download may be incomplete. &＃39;&＃39;Otherwise you can create an issue for it.&＃39;.format(self.name &＃43; &＃39;.csv.gz&＃39;))# 将文件解压缩到目录self.raw_dir下的self.name目录中self._extract_gz(gz_file_path, self.raw_path)

处理数据

假设数据已经下载到“self.raw_dir”目录下&＃xff0c;接下来就可以处理数据了。根据图上的任务&＃xff0c;分别从整图分类、节点分类和链接预测介绍。

整图分类

整图分类任务与传统机器学习任务类似&＃xff0c;整图为特征&＃xff0c;类别为标签。调用“process()”将数据集处理为 dgl.DGLGraph 对象的列表和标签张量的列表。

class QM7bDataset(DGLDataset):_url &＃61; &＃39;http://deepchem.io.s3-website-us-west-1.amazonaws.com/&＃39; \&＃39;datasets/qm7b.mat&＃39;_sha1_str &＃61; &＃39;4102c744bb9d6fd7b40ac67a300e49cd87e28392&＃39;def __init__(self, raw_dir&＃61;None, force_reload&＃61;False, verbose&＃61;False):super(QM7bDataset, self).__init__(name&＃61;&＃39;qm7b&＃39;,url&＃61;self._url,raw_dir&＃61;raw_dir,force_reload&＃61;force_reload,verbose&＃61;verbose)def process(self):mat_path &＃61; self.raw_path &＃43; &＃39;.mat&＃39;self.graphs, self.label &＃61; self._load_graph(mat_path)def _load_graph(self, filename):data &＃61; io.loadmat(filename)labels &＃61; F.tensor(data[&＃39;T&＃39;], dtype&＃61;F.data_type_dict[&＃39;float32&＃39;])feats &＃61; data[&＃39;X&＃39;]num_graphs &＃61; labels.shape[0]graphs &＃61; []for i in range(num_graphs):edge_list &＃61; feats[i].nonzero()g &＃61; dgl_graph(edge_list)g.edata[&＃39;h&＃39;] &＃61; F.tensor(feats[i][edge_list[0], edge_list[1]].reshape(-1, 1),dtype&＃61;F.data_type_dict[&＃39;float32&＃39;])graphs.append(g)return graphs, labelsdef save(self):"""save the graph list and the labels"""graph_path &＃61; os.path.join(self.save_path, &＃39;dgl_graph.bin&＃39;)save_graphs(str(graph_path), self.graphs, {&＃39;labels&＃39;: self.label})def has_cache(self):graph_path &＃61; os.path.join(self.save_path, &＃39;dgl_graph.bin&＃39;)return os.path.exists(graph_path)def load(self):graphs, label_dict &＃61; load_graphs(os.path.join(self.save_path, &＃39;dgl_graph.bin&＃39;))self.graphs &＃61; graphsself.label &＃61; label_dict[&＃39;labels&＃39;]def download(self):file_path &＃61; os.path.join(self.raw_dir, self.name &＃43; &＃39;.mat&＃39;)download(self.url, path&＃61;file_path)if not check_sha1(file_path, self._sha1_str):raise UserWarning(&＃39;File {} is downloaded but the content hash does not match.&＃39;&＃39;The repo may be outdated or download may be incomplete. &＃39;&＃39;Otherwise you can create an issue for it.&＃39;.format(self.name))&＃64;propertydef num_labels(self):return 14def __getitem__(self, idx):return self.graphs[idx], self.label[idx]def __len__(self):return len(self.graphs)


处理完数据后&＃xff0c;就可以跟传统分类任务一样使用数据了。

import dgl
import torchfrom torch.utils.data import DataLoader# 数据导入
dataset &＃61; QM7bDataset()
num_labels &＃61; dataset.num_labels# 创建collate_fn函数
def _collate_fn(batch):graphs, labels &＃61; batchg &＃61; dgl.batch(graphs)labels &＃61; torch.tensor(labels, dtype&＃61;torch.long)return g, labels# 创建 dataloaders
dataloader &＃61; DataLoader(dataset, batch_size&＃61;1, shuffle&＃61;True, collate_fn&＃61;_collate_fn)# 训练
for epoch in range(100):for g, labels in dataloader:# 用户自己的训练代码pass


节点分类

与整图分类不同&＃xff0c;节点分类通常在单个图上进行。因此数据集的划分是在图的节点集上进行。 DGL建议使用节点掩码来指定数据集的划分&＃xff0c;相当于给节点做一个标记&＃xff0c;明确是为训练节点&＃xff08;“g.ndata[‘train_mask’]”&＃xff09;、验证节点&＃xff08;“g.ndata[‘val_mask’]”&＃xff09;还是测试节点&＃xff08;“g.ndata[‘test_mask’]”&＃xff09;。 本节以内置数据集 CitationGraphDataset 为例&＃xff0c;支持’cora’, ‘citeseer’, &＃39;pubmed’三个常用的数据集&＃xff0c;DGL已经分别针对三个数据集构建了子类CoraGraphDataset、CiteseerGraphDataset和PubmedGraphDataset。

from dgl.data import DGLBuiltinDataset
from dgl.data.utils import _get_dgl_url, generate_mask_tensorclass CitationGraphDataset(DGLBuiltinDataset):_urls &＃61; {&＃39;cora_v2&＃39; : &＃39;dataset/cora_v2.zip&＃39;,&＃39;citeseer&＃39; : &＃39;dataset/citeseer.zip&＃39;,&＃39;pubmed&＃39; : &＃39;dataset/pubmed.zip&＃39;,}def __init__(self, name, raw_dir&＃61;None, force_reload&＃61;False, verbose&＃61;True):assert name.lower() in [&＃39;cora&＃39;, &＃39;citeseer&＃39;, &＃39;pubmed&＃39;]if name.lower() &＃61;&＃61; &＃39;cora&＃39;:name &＃61; &＃39;cora_v2&＃39;url &＃61; _get_dgl_url(self._urls[name])super(CitationGraphDataset, self).__init__(name,url&＃61;url,raw_dir&＃61;raw_dir,force_reload&＃61;force_reload,verbose&＃61;verbose)def process(self):# 跳过一些处理的代码# &＃61;&＃61;&＃61; 跳过数据处理 &＃61;&＃61;&＃61;# 构建图g &＃61; dgl.graph(graph)# 划分掩码g.ndata[&＃39;train_mask&＃39;] &＃61; generate_mask_tensor(train_mask)g.ndata[&＃39;val_mask&＃39;] &＃61; generate_mask_tensor(val_mask)g.ndata[&＃39;test_mask&＃39;] &＃61; generate_mask_tensor(test_mask)# 节点的标签g.ndata[&＃39;label&＃39;] &＃61; torch.tensor(labels)# 节点的特征g.ndata[&＃39;feat&＃39;] &＃61; torch.tensor(_preprocess_features(features),dtype&＃61;F.data_type_dict[&＃39;float32&＃39;])self._num_labels &＃61; onehot_labels.shape[1]self._labels &＃61; labelsself._g &＃61; gdef __getitem__(self, idx):assert idx &＃61;&＃61; 0, "这个数据集里只有一个图"return self._gdef __len__(self):return 1


由于数据集只有一个图&＃xff0c;所以需要取第0个元素“dataset[0]”&＃xff1a;

# 创建链接预测数据集示例
class KnowledgeGraphDataset(DGLBuiltinDataset):def __init__(self, name, reverse&＃61;True, raw_dir&＃61;None, force_reload&＃61;False, verbose&＃61;True):self._name &＃61; nameself.reverse &＃61; reverseurl &＃61; _get_dgl_url(&＃39;dataset/&＃39;) &＃43; &＃39;{}.tgz&＃39;.format(name)super(KnowledgeGraphDataset, self).__init__(name,url&＃61;url,raw_dir&＃61;raw_dir,force_reload&＃61;force_reload,verbose&＃61;verbose)def process(self):# 跳过一些处理的代码# &＃61;&＃61;&＃61; 跳过数据处理 &＃61;&＃61;&＃61;# 划分掩码g.edata[&＃39;train_mask&＃39;] &＃61; train_maskg.edata[&＃39;val_mask&＃39;] &＃61; val_maskg.edata[&＃39;test_mask&＃39;] &＃61; test_mask# 边类型g.edata[&＃39;etype&＃39;] &＃61; etype# 节点类型g.ndata[&＃39;ntype&＃39;] &＃61; ntypeself._g &＃61; gdef __getitem__(self, idx):assert idx &＃61;&＃61; 0, "这个数据集只有一个图"return self._gdef __len__(self):return 1


下面利用’FB15k-237’对应的子类 dgl.data.FB15k237Dataset 来做演示如何使用用于链路预测的数据集&＃xff1a;

from dgl.data import FB15k237Dataset# 导入数据
dataset &＃61; FB15k237Dataset()
graph &＃61; dataset[0]# 获取训练集掩码
train_mask &＃61; graph.edata[&＃39;train_mask&＃39;]
train_idx &＃61; torch.nonzero(train_mask).squeeze()
src, dst &＃61; graph.edges(train_idx)# 获取训练集中的边类型
rel &＃61; graph.edata[&＃39;etype&＃39;][train_idx]


保存和加载数据

DGL提供了4个函数&＃xff1a;

1. dgl.save_graphs(): 保存DGLGraph对象和标签到本地磁盘

2. dgl.load_graphs()&＃xff1a;从本地磁盘读取它们

3. dgl.data.utils.save_info(): 将数据集的有用信息(python dict对象)保存到本地磁盘

4. dgl.data.utils.load_info()和从本地磁盘读取它们

import os
from dgl import save_graphs, load_graphs
from dgl.data.utils import makedirs, save_info, load_infodef save(self):# 保存图和标签graph_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_dgl_graph.bin&＃39;)save_graphs(graph_path, self.graphs, {&＃39;labels&＃39;: self.labels})# 在Python字典里保存其他信息info_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_info.pkl&＃39;)save_info(info_path, {&＃39;num_classes&＃39;: self.num_classes})def load(self):# 从目录 &＃96;self.save_path&＃96; 里读取处理过的数据graph_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_dgl_graph.bin&＃39;)self.graphs, label_dict &＃61; load_graphs(graph_path)self.labels &＃61; label_dict[&＃39;labels&＃39;]info_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_info.pkl&＃39;)self.num_classes &＃61; load_info(info_path)[&＃39;num_classes&＃39;]def has_cache(self):# 检查在 &＃96;self.save_path&＃96; 里是否有处理过的数据文件graph_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_dgl_graph.bin&＃39;)info_path &＃61; os.path.join(self.save_path, self.mode &＃43; &＃39;_info.pkl&＃39;)return os.path.exists(graph_path) and os.path.exists(info_path)


当处理过的数据比较大时&＃xff0c;在 getitem(idx) 中处理每个数据实例是更高效的方法。

使用ogb包导入OGB数据集

OGB&＃xff08;Open Graph Benchmark&＃xff09;是一个图深度学习的基准数据集。 官方的 ogb 包提供了用于下载和处理OGB数据集到 dgl.data.DGLGraph 对象的API。

首先需要使用“pip install ogb”安装这个包&＃xff0c;接着就可以根据任务从里面加载数据集了。

图属性预测任务&＃xff08;Graph Property Prediction&＃xff09;

类的命名十分统一&＃xff0c;只需要执行“dataset &＃61; DglGraphPropPredDataset(name&＃61;‘ogbg-molhiv’)”即可得到相应的数据集&＃xff0c;然后与传统机器学习任务类似&＃xff0c;将数据处理为(graph, label)的形式。

# 载入OGB的Graph Property Prediction数据集
import dgl
import torch
from ogb.graphproppred import DglGraphPropPredDataset
from torch.utils.data import DataLoaderdef _collate_fn(batch):# 小批次是一个元组(graph, label)列表graphs &＃61; [e[0] for e in batch]g &＃61; dgl.batch(graphs)labels &＃61; [e[1] for e in batch]labels &＃61; torch.stack(labels, 0)return g, labels# 载入数据集
dataset &＃61; DglGraphPropPredDataset(name&＃61;&＃39;ogbg-molhiv&＃39;)
split_idx &＃61; dataset.get_idx_split()
# dataloader
train_loader &＃61; DataLoader(dataset[split_idx["train"]], batch_size&＃61;32, shuffle&＃61;True, collate_fn&＃61;_collate_fn)
valid_loader &＃61; DataLoader(dataset[split_idx["valid"]], batch_size&＃61;32, shuffle&＃61;False, collate_fn&＃61;_collate_fn)
test_loader &＃61; DataLoader(dataset[split_idx["test"]], batch_size&＃61;32, shuffle&＃61;False, collate_fn&＃61;_collate_fn)


节点属性预测任务&＃xff08;Node Property Prediction&＃xff09;

类似地&＃xff0c;执行“dataset &＃61; DglNodePropPredDataset(name&＃61;‘ogbn-proteins’)”即可获取数据集&＃xff0c;这种数据集只有一个图对象。

# 载入OGB的Node Property Prediction数据集
from ogb.nodeproppred import DglNodePropPredDatasetdataset &＃61; DglNodePropPredDataset(name&＃61;&＃39;ogbn-proteins&＃39;)
split_idx &＃61; dataset.get_idx_split()# there is only one graph in Node Property Prediction datasets
# 在Node Property Prediction数据集里只有一个图
g, labels &＃61; dataset[0]
# 获取划分的标签
train_label &＃61; dataset.labels[split_idx[&＃39;train&＃39;]]
valid_label &＃61; dataset.labels[split_idx[&＃39;valid&＃39;]]
test_label &＃61; dataset.labels[split_idx[&＃39;test&＃39;]]


链接属性预测任务&＃xff08;Link Property Prediction&＃xff09;

通过执行“dataset &＃61; DglLinkPropPredDataset(name&＃61;‘ogbl-ppa’)”获取数据集&＃xff0c;同样是单图。

# 载入OGB的Link Property Prediction数据集
from ogb.linkproppred import DglLinkPropPredDatasetdataset &＃61; DglLinkPropPredDataset(name&＃61;&＃39;ogbl-ppa&＃39;)
split_edge &＃61; dataset.get_edge_split()graph &＃61; dataset[0]
print(split_edge[&＃39;train&＃39;].keys())
print(split_edge[&＃39;valid&＃39;].keys())
print(split_edge[&＃39;test&＃39;].keys())