本文介绍了如何使用PyG构建内存数据集类，以 PlanetoidPubMed 数据集为例，包括InMemoryDataset的使用方法和节点预测、边预测任务的实践。详细讲解了数据完整存储、数据处理流程和节点预测网络构建。
文章目录
参考来源（datawhale组队学习）：
1. 数据完整存储于内存的数据集类
1.1 引言
对于占用内存有限的数据集，我们可以将整个数据集的数据都存储到内存里。PyG为我们提供了方便的构造数据完全存于内存的数据集类，简称为InMemory数据集类的方式，在此小节我们就将学习构造InMemory数据集类的方式。
pyG使用数据集的一般过程
- 从网络上下载数据原始文件；
- 对数据原始文件做处理，为每一个图样本生成一个
Data对象； - 对每一个
Data对象执行数据处理，使其转换成新的Data对象； - 过滤
Data对象； - 保存
Data对象到文件； - 获取
Data对象，在每一次获取Data对象时，都先对Data对象做数据变换（于是获取到的是数据变换后的Data对象）。
这么说这些步骤有些抽象，下面通过InMemoryDataset基类具体解释每个步骤。
1.2 InMemoryDataset基类简介
在PyG中，我们通过继承InMemoryDataset类来自定义一个数据可全部存储到内存的数据集类。
class InMemoryDataset ( root: Optional[str] = None, transform: Optional[Callable] = None,
pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None )
InMemoryDataset类参数说明:
-
root：字符串类型，存储数据集的文件夹的路径。该文件夹下有两个文件夹：- 一个文件夹为记录在
raw_dir，它用于存储未处理的文件，从网络上下载的数据集原始文件会被存放到这里，对应上述过程中的第一步； - 另一个文件夹记录在
processed_dir，处理后的数据被保存到这里，以后从此文件夹下加载文件即可获得Data对象，对应上述过程中的第二步。 - 注：
raw_dir和processed_dir是属性方法，我们可以自定义要使用的文件夹。
- 一个文件夹为记录在
-
transform：函数类型，一个数据转换函数，它接收一个Data对象并返回一个转换后的Data对象。此函数在每一次数据获取过程中都会被执行。获取数据的函数首先使用此函数对Data对象做转换，然后才返回数据。此函数应该用于数据增广（Data Augmentation）。该参数默认值为None，表示不对数据做转换。
-
pre_transform：函数类型，一个数据转换函数，它接收一个Data对象并返回一个转换后的Data对象。此函数在Data对象被保存到文件前调用。因此它应该用于只执行一次的数据预处理。该参数默认值为None，表示不做数据预处理，对应上述过程中的第三步。
-
pre_filter：函数类型，一个检查数据是否要保留的函数，它接收一个Data对象，返回此Data对象是否应该被包含在最终的数据集中。此函数也在Data对象被保存到文件前调用。该参数默认值为None，表示不做数据检查，保留所有的数据。
通过继承InMemoryDataset类来构造一个我们自己的数据集类，需要实现四个基本方法： -
raw_file_names()：这是一个属性方法，返回一个数据集原始文件的文件名列表，数据集原始文件应该能在raw_dir文件夹中找到，否则调用download()函数下载文件到raw_dir文件夹。 -
processed_file_names()。这是一个属性方法，返回一个存储处理过的数据的文件的文件名列表，存储处理过的数据的文件应该能在processed_dir文件夹中找到，否则调用process()函数对样本做处理，然后保存处理过的数据到processed_dir文件夹下的文件里。 -
download(): 下载数据集原始文件到raw_dir文件夹。 -
process(): 处理数据，保存处理好的数据到processed_dir文件夹下的文件。
1.3 一个简化的InMemory数据集类
基于以上介绍，我们可以通过继承InMemory类，自己构建一个简单的InMemory数据集类
import torch
from torch_geometric.data import InMemoryDataset, download_url
class MyOwnDataset(InMemoryDataset):
def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
super().__init__(root=root, transform=transform, pre_transform=pre_transform, pre_filter=pre_filter)
self.data, self.slices = torch.load(self.processed_paths[0])
#“...”也是一个对象
# https://blog.youkuaiyun.com/weixin_30578677/article/details/101792565
@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]
#数据保存
#collate()函数接收一个列表的Data对象，返回合并后的Data对象
#以及用于从合并后的Data对象重构各个原始Data对象的切片字典slices。
#最后将这个巨大的Data对象和切片字典slices保存到文件。
data, slices = self.collate(data_list)
torch.save((data, slices), self.processed_paths[0])
- 在
raw_file_names属性方法里，写上数据集原始文件有哪些。 - 在
processed_file_names属性方法里，处理过的数据要保存在哪些文件里，在此例子中只有data.pt。 - 在
download方法里，我们实现下载数据到self.raw_dir文件夹的逻辑。 - 在
process方法里，我们实现数据处理的逻辑：- 首先，我们从数据集原始文件中读取样本并生成
Data对象，所有样本的Data对象保存在列表data_list中。 - 其次，如果要对数据做过滤的话，我们执行数据过滤的过程。
- 接着，如果要对数据做处理的话，我们执行数据处理的过程。
- 最后，我们保存处理好的数据到文件。但python保存一个巨大的列表是相当慢的，我们需要先将所有
Data对象合并成一个巨大的Data对象再保存。
- 首先，我们从数据集原始文件中读取样本并生成
1.4 PlanetoidPubMed数据集类的构造
下面以公开数据集PubMed为例子，进行InMemoryDataset数据集实例分析。
PubMed数据集存储的是文章引用网络，文章对应图的结点，如果两篇文章存在引用关系（无论引用与被引），则这两篇文章对应的结点之间存在边。
该数据集来源于论文Revisiting Semi-Supervised Learning with Graph Embeddings。PyG中的Planetoid数据集类包含了数据集PubMed的使用，因此我们直接基于Planetoid类进行修改，得到PlanetoidPubMed数据集类。
import os.path as osp
import torch
from torch_geometric.data import (InMemoryDataset, download_url)
from torch_geometric.io import read_planetoid_data
class PlanetoidPubMed(InMemoryDataset):
""" 节点代表文章，边代表引文关系。
训练、验证和测试的划分通过二进制掩码给出。
参数:
root (string): 存储数据集的文件夹的路径
transform (callable, optional): 数据转换函数，每一次获取数据时被调用。
pre_transform (callable, optional): 数据转换函数，数据保存到文件前被调用。
"""
url = 'https://github.com/kimiyoung/planetoid/raw/master/data'
def __init__(self, root, transform=None, pre_transform=None):
super(PlanetoidPubMed, self).__init__(root, transform, pre_transform)
self.data, self.slices = torch.load(self.processed_paths[0])
@property
def raw_dir(self):
return osp.join(self.root, 'raw')
@property
def processed_dir(self):
return osp.join(self.root, 'processed')
@property
def raw_file_names(self):
names = ['x', 'tx', 'allx', 'y', 'ty', 'ally', 'graph', 'test.index']
return ['ind.pubmed.{}'.format(name) for name in names]
@property
def processed_file_names(self):
return 'data.pt'
def download(self):
for name in self.raw_file_names:
download_url('{}/{}'.format(self.url, name), self.raw_dir)
def process(self):
data = read_planetoid_data(self.raw_dir, 'pubmed')
data = data if self.pre_transform is None else self.pre_transform(data)
torch.save(self.collate([data]), self.processed_paths[0])
def __repr__(self):
return '{}()'.format(self.name)
该类初始化方法的参数说明见代码。代码中还实现了raw_dir()和processed_dir()两个属性方法，通过修改返回值我们就可以修改要使用的文件夹。
1.5 PlanetoidPubMed数据集类的使用
在我们生成一个PlanetoidPubMed类的对象时，程序运行流程如下：
- 首先，检查数据原始文件是否已下载：
- 检查
self.raw_dir目录下是否存在raw_file_names()属性方法返回的每个文件， - 如有文件不存在，则调用
download()方法执行原始文件下载。 self.raw_dir为osp.join(self.root, 'raw')。
- 检查
- 其次，检查数据是否经过处理：
- 首先，检查之前对数据做变换的方法：检查
self.processed_dir目录下是否存在pre_transform.pt文件：- 如果存在，意味着之前进行过数据变换，接着需要加载该文件，以获取之前所用的数据变换的方法，并检查它与当前
pre_transform参数指定的方法是否相同，- 如果不相同则会报出一个警告，“The pre_transform argument differs from the one used in ……”。
self.processed_dir为osp.join(self.root, 'processed')。
- 如果存在，意味着之前进行过数据变换，接着需要加载该文件，以获取之前所用的数据变换的方法，并检查它与当前
- 其次，检查之前的样本过滤的方法：检查
self.processed_dir目录下是否存在pre_filter.pt文件：- 如果存在，则加载该文件并获取之前所用的样本过滤的方法，并检查它与当前
pre_filter参数指定的方法是否相同，- 如果不相同则会报出一个警告，“The pre_filter argument differs from the one used in ……”。
- 如果存在，则加载该文件并获取之前所用的样本过滤的方法，并检查它与当前
- 接着，检查是否存在处理好的数据：检查
self.processed_dir目录下是否存在self.processed_file_names属性方法返回的所有文件，如有文件不存在，则需要执行以下的操作：- 调用
process()方法，进行数据处理。 - 如果
pre_transform参数不为None，则调用pre_transform()函数进行数据处理。 - 如果
pre_filter参数不为None，则进行样本过滤（此例子中不需要进行样本过滤，pre_filter参数为None）。 - 保存处理好的数据到文件，文件存储在**
processed_paths()**属性方法返回的文件路径。如果将数据保存到多个文件中，则返回的路径有多个。processed_paths()属性方法是在基类中定义的，它对self.processed_dir文件夹与processed_file_names()属性方法的返回每一个文件名做拼接，然后返回。
- 最后保存新的
pre_transform.pt文件和pre_filter.pt文件，它们分别存储当前使用的数据处理方法和样本过滤方法。
- 调用
- 首先，检查之前对数据做变换的方法：检查
2. 节点预测与边预测任务实践
2.1 节点预测
之前我们学习过由2层GATConv组成的图神经网络，现在我们重定义一个GAT图神经网络，使其能够通过参数来定义GATConv的层数，以及每一层GATConv的out_channels。我们的图神经网络定义如下：
class GAT(torch.nn.Module):
def __init__(self, num_features, hidden_channels_list, num_classes):
super(GAT, self).__init__()
torch.manual_seed(12345)
hns = [num_features] + hidden_channels_list
conv_list = []
for idx in range(len(hidden_channels_list)):
conv_list.append((GATConv(hns[idx], hns[idx+1]), 'x, edge_index -> x'))
conv_list.append(ReLU(inplace=True),)
#使用GAT和ReLU
self.convseq = Sequential('x, edge_index', conv_list)
self.linear = Linear(hidden_channels_list[-1], num_classes)
def forward(self, x, edge_index):
x = self.convseq(x, edge_index)
x = F.dropout(x, p=0.5, training=self.training)
x = self.linear(x)
return x
通过hidden_channels_list参数来设置每一层GATConv的outchannel，所以hidden_channels_list长度即为GATConv的层数。通过修改hidden_channels_list，我们就可构造出不同的图神经网络。
2.2 边预测
边预测任务，目标是预测两个节点之间是否存在边。拿到一个图数据集，我们有节点属性x，边端点edge_index。edge_index存储的便是正样本。为了构建边预测任务，我们需要生成一些负样本，即采样一些不存在边的节点对作为负样本边，正负样本数量应平衡。此外要将样本分为训练集、验证集和测试集三个集合。
PyG中为我们提供了现成的采样负样本边的方法，train_test_split_edges(data, val_ratio=0.05, test_ratio=0.1)，其
- 第一个参数为
torch_geometric.data.Data对象， - 第二参数为验证集所占比例，
- 第三个参数为测试集所占比例。
该函数将自动地采样得到负样本，并将正负样本分成训练集、验证集和测试集三个集合。它用train_pos_edge_index、train_neg_adj_mask、val_pos_edge_index、val_neg_edge_index、test_pos_edge_index和test_neg_edge_index，六个属性取代edge_index属性。
注意train_neg_adj_mask与其他属性格式不同，其实该属性在后面并没有派上用场，后面我们仍然需要进行一次训练集负样本采样。
下面我们使用Cora数据集作为例子，进行边预测任务说明。
获取数据
import os.path as osp
from torch_geometric.utils import negative_sampling
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.utils import train_test_split_edges
dataset = Planetoid('dataset', 'Cora', transform=T.NormalizeFeatures())
data = dataset[0]
data.train_mask = data.val_mask = data.test_mask = data.y = None # 不再有用
data = train_test_split_edges(data)
print(data.edge_index.shape)
# torch.Size([2, 10556])
for key in data.keys:
print(key, getattr(data, key).shape)
构造神经网络
接下来**构造神经网络**：
```python
import torch
from torch_geometric.nn import GCNConv
class Net(torch.nn.Module):
def __init__(self, in_channels, out_channels):
super(Net, self).__init__()
self.conv1 = GCNConv(in_channels, 128)
self.conv2 = GCNConv(128, out_channels)
def encode(self, x, edge_index):
x = self.conv1(x, edge_index)
x = x.relu()
return self.conv2(x, edge_index)
def decode(self, z, pos_edge_index, neg_edge_index):
edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
return (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)
def decode_all(self, z):
prob_adj = z @ z.t()
return (prob_adj > 0).nonzero(as_tuple=False).t()
模型训练
定义单个epoch训练过程
def get_link_labels(pos_edge_index, neg_edge_index):
num_links = pos_edge_index.size(1) + neg_edge_index.size(1)
link_labels = torch.zeros(num_links, dtype=torch.float)
link_labels[:pos_edge_index.size(1)] = 1.
return link_labels
def train(data, model, optimizer):
model.train()
neg_edge_index = negative_sampling(
edge_index=data.train_pos_edge_index,
num_nodes=data.num_nodes,
num_neg_samples=data.train_pos_edge_index.size(1))
optimizer.zero_grad()
z = model.encode(data.x, data.train_pos_edge_index)
link_logits = model.decode(z, data.train_pos_edge_index, neg_edge_index)
link_labels = get_link_labels(data.train_pos_edge_index, neg_edge_index).to(data.x.device)
loss = F.binary_cross_entropy_with_logits(link_logits, link_labels)
loss.backward()
optimizer.step()
return loss
定义单个epoch验证与测试过程
@torch.no_grad()
def test(data, model):
model.eval()
z = model.encode(data.x, data.train_pos_edge_index)
results = []
for prefix in ['val', 'test']:
pos_edge_index = data[f'{prefix}_pos_edge_index']
neg_edge_index = data[f'{prefix}_neg_edge_index']
link_logits = model.decode(z, pos_edge_index, neg_edge_index)
link_probs = link_logits.sigmoid()
link_labels = get_link_labels(pos_edge_index, neg_edge_index)
results.append(roc_auc_score(link_labels.cpu(), link_probs.cpu()))
return results
训练、验证与测试
def main():
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = 'Cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
ground_truth_edge_index = data.edge_index.to(device)
data.train_mask = data.val_mask = data.test_mask = data.y = None
data = train_test_split_edges(data)
data = data.to(device)
model = Net(dataset.num_features, 64).to(device)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)
best_val_auc = test_auc = 0
for epoch in range(1, 101):
loss = train(data, model, optimizer)
val_auc, tmp_test_auc = test(data, model)
if val_auc > best_val_auc:
best_val_auc = val_auc
test_auc = tmp_test_auc
print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Val: {val_auc:.4f}, '
f'Test: {test_auc:.4f}')
z = model.encode(data.x, data.train_pos_edge_index)
final_edge_index = model.decode_all(z)
if __name__ == "__main__":
main()
扫一扫
293
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
