InMemoryDataset
基类简介
在PyG中,我们通过继承InMemoryDataset
类来自定义一个数据可全部存储到内存的数据集类
class InMemoryDataset(root: Optional[str] = None, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None)
nMemoryDataset`类初始化方法参数说明:
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
文件夹中找到,否则调用process()
函数下载文件到raw_dir
文件夹。processed_file_names()
。这是一个属性方法,返回一个存储处理过的数据的文件的文件名列表,存储处理过的数据的文件应该能在processed_dir
文件夹中找到,否则调用process()
函数对样本做处理,然后保存处理过的数据到processed_dir
文件夹下的文件里。download()
: 下载数据集原始文件到raw_dir
文件夹。process()
: 处理数据,保存处理好的数据到processed_dir
文件夹下的文件。
InMemoryDataset
数据集类实例
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):
r""" 节点代表文章,边代表引用关系。
训练、验证和测试的划分通过二进制掩码给出。
参数:
root (string): 存储数据集的文件夹的路径
transform (callable, optional): 数据转换函数,每一次获取数据时被调用。
pre_transform (callable, optional): 数据转换函数,数据保存到文件前被调用。
"""
url = 'https://github.com/kimiyoung/planetoid/raw/master/data'
# url = 'https://gitee.com/rongqinchen/planetoid/raw/master/data'
# 如果github的链接不可用,请使用gitee的链接
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)
dataset = PlanetoidPubMed('dataset/PlanetoidPubMed')
print(dataset.num_classes)
print(dataset[0].num_nodes)
print(dataset[0].num_edges)
print(dataset[0].num_features)
# 3
# 19717
# 88648
# 500
节点预测与边预测任务实践
节点预测任务实践
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),)
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
,我们就可构造出不同的图神经网络。
边预测任务实践
边预测任务,目标是预测两个节点之间是否存在边。拿到一个图数据集,我们有节点属性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 # 不再有用
print(data.edge_index.shape)
# torch.Size([2, 10556])
data = train_test_split_edges(data)
for key in data.keys:
print(key, getattr(data, key).shape)
# x torch.Size([2708, 1433])
# val_pos_edge_index torch.Size([2, 263])
# test_pos_edge_index torch.Size([2, 527])
# train_pos_edge_index torch.Size([2, 8976])
# train_neg_adj_mask torch.Size([2708, 2708])
# val_neg_edge_index torch.Size([2, 263])
# test_neg_edge_index torch.Size([2, 527])
# 263 + 527 + 8976 = 9766 != 10556
# 263 + 527 + 8976/2 = 5278 = 10556/2
边预测图神经网络的构造
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()
用于做边预测的神经网络主要由两部分组成:其一是编码(encode),它与我们前面介绍的节点表征生成是一样的;其二是解码(decode),它根据边两端节点的表征生成边为真的几率(odds)。decode_all(self, z)
用于推理(inference)阶段,我们要对所有的节点对预测存在边的几率。
边预测图神经网络的训练
定义单个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
的训练过程中,都进行一次训练集负样本采样。采样到的样本数量与训练集正样本相同,但不同epoch
中采样到的样本是不同的。这样做,我们既能实现类别数量平衡,又能实现增加训练集负样本的多样性。在负样本采样时,我们传递了train_pos_edge_index
为参数,于是negative_sampling()
函数只会在训练集中不存在边的节点对中采样。get_link_labels()
函数用于生成完整训练集的标签。
定义单个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()