图神经网络_04-数据完整存储与内存的数据集类+节点预测与边预测任务实践

数据完整存储与内存的数据集类

使用数据集的一般过程

PyG使用数据的一般过程：

1. 从网络上下载数据原始文件；
2. 对原始数据文件做处理，为每一个图样本生成一个Data对象；
3. 对每一个Data对象执行数据处理，使其转换成新的Data对象；
4. 过滤Data对象；
5. 保存Data对象到文件；
6. 获取Data对象，在每一次获取Data对象时，都先对Data对象做数据变换，所以获取到的是数据变换后的Data对象）。

实际中并非需要严格执行每一个步骤

占用内存较小的数据集，我们可以将整个数据集的数据都存储到内存里。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对象。
• transform：函数类型，一个数据转换函数，接收Data对象并返回转换后的Data对象。此函数在每一次数据获取过程中都会被执行。获取数据的函数首先使用此函数对Data对象做转换，然后才返回数据。此函数应该用于数据增广（Data Augmentation）。默认为None，不对数据做转换。
• pre_transform：函数类型，一个数据转换函数，它接收Data对象并返回转换后的Data对象。此函数在Data对象被保存到文件前调用。因此它应该用于只执行一次的数据预处理。该参数默认值为None，表示不做数据预处理。
• pre_filter：函数类型，检查数据是否要保留的函数，它接收Data对象，返回此Data对象是否应该被包含在最终的数据集中。此函数也在[Data]对象被保存到文件前调用。默认为None，表示不做数据检查，保留所有的数据。

InMemoryDataset官方文档

通过继承[InMemoryDataset]类来构造一个我们自己的数据集类，我们需要实现四个基本方法：

• raw_file_names()：属性方法，返回一个数据集原始文件的文件名列表，若在raw_dir文件夹中找不到数据集原始文件，则调用process()函数下载文件到raw_dir文件夹。
• processed_file_names()。属性方法，返回一个存储处理过的数据的文件的文件名列表，若在processed_dir文件夹中找不到存储处理过的数据的文件，则调用process()函数对样本做处理，然后保存处理过的数据到processed_dir文件夹下的文件里。
• download(): 下载数据集原始文件到raw_dir文件夹。
• process(): 处理数据，保存处理好的数据到processed_dir文件夹下的文件。

通过继承自InMemoryDataset自定义的数据集类

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])

    @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]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

InMemoryDataset数据集类实例

PubMed数据集存储的是文章引用网络，文章对应图的结点，如果两篇文章存在引用关系，则这两篇文章对应的结点之间存在边。PyG中的Planetoid数据集类包含了数据集PubMed的使用，因此我们直接基于Planetoid类进行修改，得到PlanetoidPubMed数据集类。

• PubMed数据集来源于论文Revisiting Semi-Supervised Learning with Graph Embeddings。

构造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):
    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)

查看这个数据集

dataset = PlanetoidPubMed('dataset/PlanetoidPubMed')
print(dataset.num_classes)
print(dataset[0].num_nodes)
print(dataset[0].num_edges)
print(dataset[0].num_features)

Processing...
Done!
3
19717
88648
500

节点预测与边预测任务实践

1.节点预测

之前设计过2层GATConv组成的图神经网络，现重定义一个GAT网络，使其能够通过参数来定义GATConv的层数，以及每一层GATConv的out_channels。

import os.path as osp
import torch
from torch_geometric.data import (InMemoryDataset, download_url)
from torch_geometric.io import read_planetoid_data
from torch_geometric.transforms import NormalizeFeatures
from torch_geometric.nn import Sequential, GATConv
from torch.nn import ReLU, Linear
import torch.nn.functional as F

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

定义网络训练和测试

def train(model, data, criterion, optimizer):
    model.train()
    optimizer.zero_grad()  # Clear gradients.
    out = model(data.x, data.edge_index)  # Perform a single forward pass.
    loss = criterion(out[data.train_mask],
                     data.y[data.train_mask])  # Compute the loss solely based on the training nodes.
    loss.backward()  # Derive gradients.
    optimizer.step()  # Update parameters based on gradients.
    return loss


def test(model, data):
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)  # Use the class with highest probability.
    test_correct = pred[data.test_mask] == data.y[data.test_mask]  # Check against ground-truth labels.
    test_acc = int(test_correct.sum()) / int(data.test_mask.sum())  # Derive ratio of correct predictions.
    return test_acc

我们来训练一下这个网络看看效果

def run_gat():
    dataset = PlanetoidPubMed(root='dataset/PlanetoidPubMed/', transform=NormalizeFeatures())
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    data = dataset[0].to(device)
    hidden_channels_list = [100, 200]

    model = GAT(num_features=dataset.num_features,
                hidden_channels_list=hidden_channels_list,
                num_classes=dataset.num_classes).to(device)
    print(model)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
    criterion = torch.nn.CrossEntropyLoss()

    # init_lr = optimizer.param_groups[0]['lr']
    for epoch in range(1, 501):
        # adjust_learning_rate(optimizer, epoch, init_lr)
        # lr = optimizer.param_groups[0]['lr']
        # print(epoch, lr)

        loss = train(model, data, criterion, optimizer)
        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')

    test_accuray = test(model, data)
    print(f'Test Accuracy: {test_accuray:.4f}')

GAT(
  (convseq): Sequential(
    (0): GATConv(500, 100, heads=1)
    (1): ReLU(inplace=True)
    (2): GATConv(100, 200, heads=1)
    (3): ReLU(inplace=True)
  )
  (linear): Linear(in_features=200, out_features=3, bias=True)
)
Epoch: 001, Loss: 1.0990
Epoch: 002, Loss: 1.0965
Epoch: 003, Loss: 1.0926
Epoch: 004, Loss: 1.0850
Epoch: 005, Loss: 1.0739
Epoch: 006, Loss: 1.0524
Epoch: 007, Loss: 1.0260
Epoch: 008, Loss: 0.9887
Epoch: 009, Loss: 0.9284
Epoch: 010, Loss: 0.8804
...
Epoch: 496, Loss: 0.0060
Epoch: 497, Loss: 0.0069
Epoch: 498, Loss: 0.0064
Epoch: 499, Loss: 0.0076
Epoch: 500, Loss: 0.0060
Test Accuracy: 0.7630

2.边预测

边预测任务，目标是预测两个节点之间是否存在边。
对于一个图数据集，我们有节点属性x，边端点edge_index。edge_index存储的就是正样本。为了实现边预测任务，我们需要生成部分负样本，即采样一些不存在边的节点对作为负样本边，正负样本数量应相当, 以便保持平衡。数据集拆分为训练集、验证集和测试集三个集合。

PyG提供了采样负样本边的方法
train_test_split_edges(data, val_ratio=0.05, test_ratio=0.1)

• data参数为torch_geometric.data.Data对象，
• val_ratio参数为验证集所占比例，
• test_ratio参数为测试集所占比例。

该函数将自动采样得到负样本，并将正负样本拆分成训练集、验证集和测试集三个集合。edge_index属性转为train_pos_edge_index、train_neg_adj_mask、val_pos_edge_index、val_neg_edge_index、test_pos_edge_index和test_neg_edge_index六个属性。

注:train_neg_adj_mask与其他属性格式不同，其实该属性在后面并没有派上用场，后面我们仍然需要进行一次训练集负样本采样。

下面我们使用Cora数据集作为例子，进行边预测任务说明。

from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.utils import train_test_split_edges

# if __name__ == '__main__':

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)

data = train_test_split_edges(data)

for key in data.keys:
    print(key, getattr(data, key).shape)

torch.Size([2, 10556])
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])

训练集、验证集和测试集中正样本边的数量之和不等于原始边的数量。因为Cora图是无向图，在统计原始边数量时，统计每一条边正向与反向各计了一次，训练集也包含边的正向与反向，但验证集与测试集都只包含了边的一个方向。因为训练集用于训练时一条边的两个端点要互传信息，只考虑一个方向的话，表示只能由一个端点传信息给另一个端点，而验证集与测试集的边用于衡量检验边预测的准确性，只需考虑一个方向的边即可。

边预测图神经网络的构造

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

定义测试方法

@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'
    dataset = Planetoid('dataset', 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, 301):
        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 Accuracy: {val_auc:.4f}, '
              f'Test Accuracy: {test_auc:.4f}')

    z = model.encode(data.x, data.train_pos_edge_index)
    final_edge_index = model.decode_all(z)

训练结果

Epoch: 001, Loss: 0.6930, Val Accuracy: 0.6852, Test Accuracy: 0.6838
Epoch: 002, Loss: 0.6816, Val Accuracy: 0.6784, Test Accuracy: 0.6838
Epoch: 003, Loss: 0.7089, Val Accuracy: 0.6838, Test Accuracy: 0.6838
Epoch: 004, Loss: 0.6767, Val Accuracy: 0.6979, Test Accuracy: 0.6983
Epoch: 005, Loss: 0.6843, Val Accuracy: 0.7087, Test Accuracy: 0.7065
...
Epoch: 096, Loss: 0.4456, Val Accuracy: 0.9141, Test Accuracy: 0.9084
Epoch: 097, Loss: 0.4450, Val Accuracy: 0.9146, Test Accuracy: 0.9084
Epoch: 098, Loss: 0.4357, Val Accuracy: 0.9148, Test Accuracy: 0.9084
Epoch: 099, Loss: 0.4415, Val Accuracy: 0.9137, Test Accuracy: 0.9084
Epoch: 100, Loss: 0.4431, Val Accuracy: 0.9129, Test Accuracy: 0.9084