（Datawhale）图神经网络Task04：数据完全存储于内存的数据集类+节点预测与边预测实践

数据完整存储于内存的数据集类构造

PyG规定的使用数据的一般过程

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

InMemoryDataset基类

PyG提供了很多官方数据集，但是当需要构建自己的数据集的时候，就需要使用到dataset基类。PyG提供了两个构建数据集的基类：torch_geometric.data.Dataset和torch_geometric.data.InMemoryDataset，其中torch_geometric.data.InMemoryDataset继承了torch_geometric.data.Dataset，表示是否将整个数据集加载到内存中。

• InMemoryDataset基类：class InMemoryDataset(root, transform, pre_transform, pre_filter)，其中：

  • root：存储数据集的文件夹的路径。该文件夹下有两个文件夹：一个文件夹为raw_dir，用于存储未处理的文件；另一个文件夹为processed_dir，用于存储处理后的数据，以后可以直接从文件夹下加载数据获得Data对象；
  • transform：数据转换函数，接收一个Data对象并返回一个转换后的Data对象，在每一次数据获取过程中都会被执行’
  • pre_transform：数据转换函数，接收一个Data对象并返回一个转换后的Data对象，在Data对象被保存到文件前调用；
  • pre_filter：检查数据是否要保留的函数，它接收一个Data对象，返回此Data对象是否应该被包含在最终的数据集中，在Data对象被保存到文件前调用。

• 通过继承InMemoryDataset类构造自己的数据集类，需要实现以下四个基本方法：

  • raw_file_names()：该方法返回数据集原始文件的文件名列表，原始文件应存在于raw_dir文件夹，否则调用download()函数下载文件到raw_dir文件夹；
  • processed_file_names()：该方法返回处理过的数据文件的文件名列表，处理文件应存在于processed_dir文件夹，否则调用process()函数对原始文件进行处理并保存到相应文件夹；
  • download()：下载数据集原始文件到raw_dir文件夹；
  • process()：处理数据并将保存处理好的数据到processed_dir文件夹下的文件。

• 上述构造过程的关键在于process()方法，该方法读取样本并生成Data对象，并将所有样本的Data对象保存在到列表data_list。由于保存python列表效率较低，因此在保存前通过torch_geometric.data.InMemoryDataset.collate()方法将列表合并成一个大的Data对象，该过程将所有对象连接成一个大数据对象，同时，返回一个切片字典以便从该对象重建单个小的对象。将合并的Data对象和切片字典slices保存到文件，并可加载到self.data和 self.slices属性。总之，process过程一共分四步：

  • 加载数据并创建列表data_list；
  • 进行数据处理过程，包括pre_filter、pre_transform；
  • 调用collate()函数；
  • 存储本地，包括data、slices。

  import torch
  from torch_geometric.data import InMemoryDataset, download_url
  
  class MyOwnDataset(InMemoryDataset):
      def __init__(self, root, transform=None, pre_transform=None):
          super(MyOwnDataset, self).__init__(root, transform, pre_transform)
          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的数据集构造实例

以Planetoid数据集类中的PubMed为例。在生成PlanetoidPubMed类的对象时，程序运行流程如下：

• 首先，检查数据原始文件是否已下载：
  • 检查self.raw_dir目录下是否存在raw_file_names()属性方法返回的每个文件；
  • 如有文件不存在，则调用download()方法执行原始文件下载。
• 其次，检查数据是否经过处理：
  • 检查self.processed_dir目录下是否存在pre_transform.pt文件。如果存在，意味着之前进行过数据变换，接着需要加载该文件，以获取之前所用的数据变换的方法，并检查它与当前pre_transform参数指定的方法是否相同；
  • 检查self.processed_dir目录下是否存在pre_filter.pt文件。如果存在，则加载该文件并获取之前所用的样本过滤的方法，并检查它与当前pre_filter参数指定的方法是否相同。
  • 检查self.processed_dir目录下是否存在self.processed_file_names属性方法返回的所有文件，如有文件不存在，则需要执行以下的操作：
    • 调用process()方法，进行数据处理；
    • 如果pre_transform参数不为None，则调用pre_transform()函数进行数据处理；
    • 如果pre_filter参数不为None，则调用pre_filter()函数进行样本过滤；
    • 保存处理好的数据到文件，文件存储在processed_paths()属性方法返回的文件路径，processed_paths()属性对self.processed_dir文件夹与processed_file_names()属性方法的返回每一个文件名做拼接。

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://gitee.com/rongqinchen/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)  # 3
print(dataset[0].num_nodes)  # 19717
print(dataset[0].num_edges)  # 88648
print(dataset[0].num_features)  # 500

PlanetoidPubMed继承InMemoryDataset，InMemoryDataset继承Dataset。

数据集下载时报错out = pickle.load(f, encoding='latin1') _pickle.UnpicklingError: invalid load key, '<'.：

• 报错原因：planetoid.py下载的数据格式为网页，不是真实的数据（？），反正就是直接下载服务器的数据文件时出现了问题。
• 解决办法：手动下载数据集，将其放置到self.raw_dir文件夹下，然后再此执行程序。此时由于dataset/PlanetoidPubMed/raw文件夹下已经存在原数据集，程序检测到后会跳过下载步骤，直接对该数据集进行处理。

通过from torch_geometric.datasets import Planetoid dataset = Planetoid(root='dataset', name='PubMed')，我们可以得到与上述自建数据集类相同的结果。从下图可以看到，Planetoid也是继承了InMemoryDataset，通过raw_file_names()、processed_file_names()、download()、process()四个基本方法构造了"Cora", “CiteSeer” 和 "PubMed"三个论文引用数据集。

基于节点表征的节点预测和边预测实践

使用上文的PlanetoidPubMed数据集类进行节点预测与边预测任务。

节点预测

import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv, Sequential
from torch_geometric.transforms import NormalizeFeatures
from torch.nn import Linear, ReLU


dataset = PlanetoidPubMed(root='../../dataset/PlanetoidPubMed/', transform=NormalizeFeatures())
print('dataset.num_features:', dataset.num_features)  # 500
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = dataset[0].to(device)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss

def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    test_correct = pred[data.test_mask] == data.y[data.test_mask]
    test_acc = int(test_correct.sum()) / int(data.test_mask.sum())
    return test_acc


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  # [500, 200, 100]
        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

model = GAT(num_features=dataset.num_features, hidden_channels_list=[200, 100], num_classes=dataset.num_classes).to(device)
print(model)
# GAT(
#   (convseq): Sequential(
#     (0): GATConv(500, 200, heads=1)
#     (1): ReLU(inplace=True)
#     (2): GATConv(200, 100, heads=1)
#     (3): ReLU(inplace=True)
#   )
#   (linear): Linear(in_features=100, out_features=3, bias=True)
# )
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(1, 201):
    loss = train()
    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}')


test_acc = test()
print(f'Test Accuracy: {test_acc:.4f}')
# Test Accuracy: 0.7560

• 数据集的训练集、验证集、测试集划分：

  • 在read_planetoid_data方法中，训练集、验证集、测试集的一般划分规则为训练集大小为y.size()，验证集大小为500，测试集大小取决于test.index文件，见下图：

    

  • 在from torch_geometric.datasets import Planetoid类方法中，训练集、验证集、测试集的划分方法分三种：public、full、random，其中public方法即上述一般方法，full、random方法如下图：

  • 

• torch_geometric.nn.Sequential类。由于 GNN 运算符接受多个输入参数，Sequential需要全局输入参数和独立运算的函数头定义。在下述代码中，'x, edge_index'定义了模型全局输入参数， 'x, edge_index -> x'定义了GCNConv的函数头，即输入参数和返回类型。如果省略，中间模块将对其前一个模块的输出进行操作，这样可能导致错误，例如对于第二个GCNConv，其接收x, edge_index作为图卷积网络的输入，如果在Sequential未定义函数头'x, edge_index -> x'，那么GCNConv将上一层ReLU的输出作为输入，引起错误。

  from torch.nn import Linear, ReLU
  from torch_geometric.nn import Sequential, GCNConv
  
  model = Sequential('x, edge_index',[
      (GCNConv(in_channels, 64), 'x, edge_index -> x'),
      ReLU(inplace=True),
      (GCNConv(64, 64), 'x, edge_index -> x'),
      ReLU(inplace=True),
      Linear(64, out_channels),
  ])
  

• 结合全局输入参数和独立运算的函数头，Sequential可以方便的实现复杂图神经网络的定义：

  from torch.nn import Linear, ReLU, Dropout
  from torch_geometric.nn import Sequential, GCNConv, JumpingKnowledge
  from torch_geometric.nn import global_mean_pool
  
  model = Sequential('x, edge_index, batch', [
      (Dropout(p=0.5), 'x -> x'),
      (GCNConv(dataset.num_features, 64), 'x, edge_index -> x1'),
      ReLU(inplace=True),
      (GCNConv(64, 64), 'x1, edge_index -> x2'),
      ReLU(inplace=True),
      (lambda x1, x2: [x1, x2], 'x1, x2 -> xs'),
      (JumpingKnowledge("cat", 64, num_layers=2), 'xs -> x'),
      (global_mean_pool, 'x, batch -> x'),
      Linear(2 * 64, dataset.num_classes),
  ])
  

边预测

边预测任务：预测两个节点之间是否存在边。

Cora数据集预处理

拿到一个图数据集，有节点属性x，边端点edge_index。其中，edge_index存储的便是正样本，在神经网络的训练中，需要生成一些负样本，即采样一些不存在边的节点对作为负样本边，正负样本数量应平衡。此外要将所有正负样本分为训练集、验证集和测试集三个集合。

PyG中提供了现成的负样本边生成方法，train_test_split_edges(data, val_ratio=0.05, test_ratio=0.1)。

• 该函数将自动地采样得到负样本，并将正负样本分成训练集、验证集和测试集三个集合，即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与其他属性格式不同，其实该属性在后面并没有派上用场，后面我们仍然需要进行一次训练集负样本采样。

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

可以看到，训练集、验证集和测试集中正样本边的数量之和不等于原始边的数量。这是因为，Cora是无向图，在统计原始边数量时，每一条边的正向与反向各统计了一次，而划分后的验证集与测试集都只包含了边的一个方向，仅训练集也包含边的正向与反向。

**为什么训练集要包含边的正向与反向，而验证集与测试集都只包含了边的一个方向？**这是因为，训练集用于训练，训练时一条边的两个端点要互传信息，只考虑一个方向的话，只能由一个端点传信息给另一个端点，而验证集与测试集的边用于衡量检验边预测的准确性，只需考虑一个方向的边即可。

边预测图神经网络

用于做边预测的神经网络主要由两部分组成：

• 其一是编码（encode），它与我们前面介绍的节点表征生成是一样的；
• 其二是解码（decode），它根据边两端节点的表征生成边为真的几率（odds）。
• decode_all(self, z)用于推理（inference）阶段，我们要对所有的节点对预测存在边的几率。

import os.path as osp
import torch
import torch.nn.functional as F
import torch_geometric.transforms as T
from sklearn.metrics import roc_auc_score
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling, train_test_split_edges


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


def get_link_labels(pos_edge_index, neg_edge_index):
    # 将正样本（存在边）和负样本（无边）的标签拼接，即torch([1.,1.,1.···0.,0.,0.])
    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))  # 采样数与正样本相等

    # # 若训练集负样本与测试集或验证集的正样本存在交集，会输出“wrong”提式，该部分代码可注释掉
    # train_neg_edge_set = set(map(tuple, neg_edge_index.T.tolist()))
    # val_pos_edge_set = set(map(tuple, data.val_pos_edge_index.T.tolist()))
    # test_pos_edge_set = set(map(tuple, data.test_pos_edge_index.T.tolist()))
    # if (len(train_neg_edge_set & val_pos_edge_set) > 0) or (len(train_neg_edge_set & test_pos_edge_set) > 0):
    #     print('wrong!')

    optimizer.zero_grad()
    z = model.encode(data.x, data.train_pos_edge_index)  # [num_nodes, out_channels],i.e.[2708, 64]
    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 = 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)
    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)  # z = model.encode(data.x, data.edge_index)
    final_edge_index = model.decode_all(z)
    
    return final_edge_index


if __name__ == "__main__":
    final_edge_index = main()
# final_edge_index.shape: torch.Size([2, 3338592])

# 经统计，train_pos_edge_index中的8967条边包含2623个节点，Cora的总结点数为2708。为此，上述代码中106行对所有节点对是否存在边的预测修改为z = model.encode(data.x, data.edge_index)。

在图中，存在边的节点对的数量往往少于不存在边的节点对的数量。在每一个epoch的训练过程中，都进行一次训练集负样本采样。采样到的样本数量与训练集正样本数量相同，但不同epoch中采样到的样本是不同的。这样做，既能实现类别数量平衡，又能实现增加训练集负样本的多样性。每一epoch训练过程中的负样本采样通过negative_sampling方法实现：

def negative_sampling(edge_index, num_nodes=None, num_neg_samples=None,
                      method="sparse", force_undirected=False):
    # 根据给出的edge_index进行负采样，如给定用于训练的正样本train_pos_edge_index，函数只会对包含于训练集节点对不存在的边负采样
    # edge_index：提供采样依赖的边索引
    # num_nodes：节点数量，默认为“None”，无需提供准确数值，函数中通过num_nodes = maybe_num_nodes(edge_index, num_nodes)计算可能的节点数，以辅助负采样
    # num_neg_samples：负采样数量，默认为“None”，此时采样数与提供的正样本edge_index边数一致
    # method："sparse"和"dense"两种方式，"sparse"适用所有大小的图，选择"dense"且适用时会有更快的采样速度。
    # force_undirected：若为True，负采样是无向边。

以data.train_pos_edge_index为实际参数来进行训练集负样本采样，但这样采样得到的负样本可能包含一些验证集的正样本与测试集的正样本，即可能将真实的正样本标记为负样本，由此会产生冲突。但该问题不大：每次训练都会重新采样，经统计100次采样中61次出现冲突，冲突的个数每次仅为1、2个，占比极小，影响不大，由此产生的微小误差模型应该是可以修正的。

参考

1. Datawhale图深度学习