DataWhale图网络学习（四）数据完全存于内存的节点预测与边预测实践

1 数据完全存于内存的数据集类

1.1 InMemoryDataset基类简介

在PyG中，我们通过继承InMemoryDataset类来自定义一个数据可全部存储到内存的数据集类：

class InMemoryDataset(root: Optional[str] = None, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None)

• 根文件夹（root），它指示数据集应该被保存在哪里。在根目录下至少有两个文件夹：1）一个文件夹为 raw_dir，它用于存储未处理的文件，从网络上下载的数据集文件会被存放到这里；2）另一个文件夹为 processed_dir，处理后的数据集被保存到这里。
• transform 函数接受Data对象为参数，对其转换后返回。此函数在每一次数据访问时被调用，所以它应该用于数据增广（Data Augmentation）。
• pre_transform 函数接受 Data对象为参数，对其转换后返回。此函数在样本 Data对象保存到文件前调用，所以它最好用于只需要做一次的大量预计算。
• pre_filter 函数可以在保存前手动过滤掉数据对象。该函数的一个用例是，过滤样本类别。

我们如果想要自定义一个InMemoryDataset类，我们需要实现如下四种方法：

• raw_file_names()这是一个属性方法，返回一个文件名列表，文件应该能在raw_dir文件夹中找到，否则调用download()函数下载文件到raw_dir文件夹。
• processed_file_names()。这是一个属性方法，返回一个文件名列表，文件应该能在processed_dir文件夹中找到，否则调用process()函数对样本做预处理然后保存到processed_dir文件夹。
• download(): 将原始数据文件下载到raw_dir文件夹。
• process(): 对样本做预处理然后保存到processed_dir文件夹。

下面我们就来看一下如何自定义一个数据集类吧:

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, split="public", num_train_per_class=20,
                 num_val=500, num_test=1000, transform=None,
                 pre_transform=None):

        super(PlanetoidPubMed, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

        self.split = split
        assert self.split in ['public', 'full', 'random']

        if split == 'full':
            data = self.get(0)
            data.train_mask.fill_(True)
            data.train_mask[data.val_mask | data.test_mask] = False
            self.data, self.slices = self.collate([data])

        elif split == 'random':
            data = self.get(0)
            data.train_mask.fill_(False)
            for c in range(self.num_classes):
                idx = (data.y == c).nonzero(as_tuple=False).view(-1)
                idx = idx[torch.randperm(idx.size(0))[:num_train_per_class]]
                data.train_mask[idx] = True

            remaining = (~data.train_mask).nonzero(as_tuple=False).view(-1)
            remaining = remaining[torch.randperm(remaining.size(0))]

            data.val_mask.fill_(False)
            data.val_mask[remaining[:num_val]] = True

            data.test_mask.fill_(False)
            data.test_mask[remaining[num_val:num_val + num_test]] = True

            self.data, self.slices = self.collate([data])

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

这里我们使用的是PubMed数据集，该数据集存储的是文章引用网络，文章对应图的节点，如果两篇文章存在引用关系（无论引用与被引），则这两篇文章对应的结点之间存在边。
当数据集类被创建时，其运行流程如下：

• 1、首先检查数据原始文件是否已下载：
  • 检查self.raw_dir目录下是否存在raw_file_names()属性方法返回的每个文件，
  • 如有文件不存在，则调用download()方法执行原始文件下载。
  • 其中self.raw_dir为self.root+ 'raw'。
• 2、其次检查数据是否经过处理：
  • 首先检查之前对数据做变换的方法：检查self.processed_dir目录下是否存在pre_transform.pt文件：如果存在，意味着之前进行过数据变换，那么加载该文件以获取之前所用的数据变换的方法，并检查它与当前pre_transform参数指定的方法是否相同；如果不相同则会报出一个警告：“The pre_transform argument differs from the one used in ……”。
  • 接着检查之前的样本过滤的方法：检查self.processed_dir目录下是否存在pre_filter.pt文件，如果存在，意味着之前进行过样本过滤，那么加载该文件以获取之前所用的样本过滤方法，并检查它与当前pre_filter参数指定的方法是否相同，如果不相同则会报出一个警告：“The pre_filter argument differs from the one used in ……”。
  • 其中self.processed_dir为self.root+'processed'。
  • 接着检查是否存在处理好的数据：检查self.processed_dir目录下是否存在self.processed_paths方法返回的所有文件（这里是data.pt），如有文件不存在，意味着不存在已经处理好的样本的文件，如需执行以下的操作：
    • 调用process方法，进行数据处理。
    • 如果pre_transform参数不为None，则调用pre_transform方法进行数据处理。
    • 如果pre_filter参数不为None，则进行样本过滤（此例子中不需要进行样本过滤，pre_filter参数始终为None）。
    • 保存处理好的数据到文件，文件存储在processed_paths()属性方法返回的路径。如果将数据保存到多个文件中，则返回的路径有多个。这些路径都在self.processed_dir目录下，以processed_file_names()属性方法的返回值为文件名。
    • 最后保存新的pre_transform.pt文件和pre_filter.pt文件，其中分别存储当前使用的数据处理方法和样本过滤方法。

最后对数据集方法进行总结：首先检查是否有原始数据和处理后的数据，没有就下载和处理，并保存，这样就可以形成一个数据集了。

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

3
19717
88648
500

该数据集存在三个分类任务，共19717个节点，88648条边，节点特征维度为500。

2 节点预测实践

在第三部分的学习中，我们已经学习过如何定义一个GAT网络，此时我们同样通过GATConv函数定义一个GAT网络。
首先导入需要的包：

import os.path as osp

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

接下来我们定义GAT模型，这里我们将GATConv作为隐藏层，将filter的数量作为参数进行显式定义，输入维度为num_features，每一层的输出维度定义在hidden_channels_list中，并通过hns加载到每一层中。

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

---

注：
这里我们使用了Sequential(args: str, modules: List[Union[Tuple[Callable, str], Callable]])函数用于定义顺序GNN模型。由于GNN操作符接受多个输入参数，因此torch_geometry .nn.Sequential需要全局输入参数和单个操作符的函数头定义。例如：

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

其中x, edge_index定义了model的输入参数，x, edge_index -> x定义了GCNConv的函数头，即输入参数和返回类型。

---

再接下来我们定义训练函数train()、评估函数value()和测试函数test()：

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

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

def test():
    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

然后利用前边已经定义的数据集类，以及制作好数据集的PubMed数据集，进行训练和学习，我们先加载全部数据集到内存中：

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

dataset.num_features: 500

然后我们设置模型的优化器、损失函数等：

model = GAT(num_features=dataset.num_features, hidden_channels_list=[200, 100], 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()

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

接下来就是训练了，这里我们训练200epoch：

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

Epoch: 001, Loss: 1.0986, Valid Loss: 1.0991, Valid Accuracy: 0.1960
Epoch: 002, Loss: 1.0943, Valid Loss: 1.0946, Valid Accuracy: 0.4400
Epoch: 003, Loss: 1.0891, Valid Loss: 1.0886, Valid Accuracy: 0.4700
Epoch: 004, Loss: 1.0745, Valid Loss: 1.0773, Valid Accuracy: 0.5140
Epoch: 005, Loss: 1.0528, Valid Loss: 1.0614, Valid Accuracy: 0.5120
Epoch: 006, Loss: 1.0186, Valid Loss: 1.0360, Valid Accuracy: 0.5220
Epoch: 007, Loss: 0.9746, Valid Loss: 1.0001, Valid Accuracy: 0.5440
Epoch: 008, Loss: 0.9236, Valid Loss: 0.9581, Valid Accuracy: 0.6120
Epoch: 009, Loss: 0.8754, Valid Loss: 0.9076, Valid Accuracy: 0.6540
Epoch: 010, Loss: 0.7773, Valid Loss: 0.8607, Valid Accuracy: 0.6620

Epoch: 191, Loss: 0.0082, Valid Loss: 0.7990, Valid Accuracy: 0.7780
Epoch: 192, Loss: 0.0074, Valid Loss: 0.8363, Valid Accuracy: 0.7700
Epoch: 193, Loss: 0.0079, Valid Loss: 0.9154, Valid Accuracy: 0.7420
Epoch: 194, Loss: 0.0185, Valid Loss: 0.8164, Valid Accuracy: 0.7640
Epoch: 195, Loss: 0.0064, Valid Loss: 0.7783, Valid Accuracy: 0.7800
Epoch: 196, Loss: 0.0070, Valid Loss: 0.7717, Valid Accuracy: 0.7700
Epoch: 197, Loss: 0.0103, Valid Loss: 0.7947, Valid Accuracy: 0.7760
Epoch: 198, Loss: 0.0068, Valid Loss: 0.8582, Valid Accuracy: 0.7500
Epoch: 199, Loss: 0.0073, Valid Loss: 0.9108, Valid Accuracy: 0.7340
Epoch: 200, Loss: 0.0059, Valid Loss: 0.9525, Valid Accuracy: 0.7280

从训练损失、评估损失和评估准确率我们可以看到，前10个epoch的学习是正常的，到最后10个已经过拟合了，我们利用测试集看一下测试效果吧

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

Test Accuracy: 0.7460

可以看到测试准确率只有74.6%，比较低，也说明模型过拟合，产生这一结果的原因可能是训练数据集中节点数量很少，我们看到总共有19717个节点，只有60个作为训练，500个验证，以及1000个测试。

Number of nodes: 19717
Number of training nodes: 60
Number of valid nodes: 500
Number of testing nodes: 1000

3 边预测实践

在边预测任务中，我么使用Core数据集，在该数据集中有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属性取代edge_index属性，但是注意train_neg_adj_mask与其他属性格式不同，但是该属性在后面并没有派上用场，后面我们仍然需要进行一次负样本训练集采样。

我们对数据进行采样和划分：

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

dataset = 'Cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
data.train_mask = data.val_mask = data.test_mask = data.y = None
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])
这里也可以看到train_neg_adj_mask不同于其他属性

我们观察到三个集合中正样本边的数量之和不等于原始边的数量。这是因为原始边的数量统计的是双向边的数量，在验证集正样本边和测试集正样本边中只需对一个方向的边做预测精度的衡量，对另一个方向的预测精度衡量属于重复，但在训练集还是保留双向的边（也可以不保留）。

接下来构建网络，该网络可以分为编码encode和解码decode两部分，编码部分与前边生成节点表征是相同的，解码部分是根据节点表征生成两端点之间有边连接为真的几率odds。decode_all(self, z)用于推断（inference）阶段，我们要对输入节点所有的节点对之间预测其存在边连接的几率。

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

Net(
(conv1): GCNConv(1433, 128)
(conv2): GCNConv(128, 64)
)

接下来就要定义模型训练部分，但是从前边对数据的分析，我们可以看到正负样本的数量并不相等，我们采用过采样方法是的采样与正样本数量一样多的负样本，这样也可以扩大数据集，丰富负样本。在负样本采样negative_sampling时，我们传递了train_pos_edge_index为参数，于是negative_sampling函数只会在训练集中不存在边的结点对中采样。也就是说我们并没有使用train_neg_adj_mask，而是进行了重新采样。

---

注：
这里我们使用negative_sampling(edge_index, num_nodes=None, num_neg_samples=None, method='sparse', force_undirected=False)进行负样本采样。

• edge_index：边的索引。
• num_nodes：网络中节点的数量。
• num_neg_samples：返回的负样本的(大致)数目。如果设置为None，将尝试为每条正边返回一个负边。
• method：抽取负样本所用的方法，例如sparse或者 dense。sparse适用于任何大小的图，而 dense可以执行更快的真负值检查。
• force_undirected：如果设置为True，采样的负边将是无向的（默认：False）。

---

get_link_labels函数用于生成完整训练集的标签。

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

接下来我们实例化模型，并设置优化器

model = Net(dataset.num_features, 64).to(device)
optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)

接下来我们训练模型，进行了100epoch：

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)

tensor([[ 0, 0, 0, …, 2707, 2707, 2707],
[ 0, 2, 4, …, 2705, 2706, 2707]])

参考：

• DataWhale开源内容
• torch_geometric.utils官方文档
• torch_geometric.nn官方文档