数据完整存储于内存的数据集类+节点预测与边预测任务实践

一、完全基于内存的数据集类

InMemoryDataset基类简介

import torch
from torch_geometric.data import InMemoryDataset, download_url

class InMemoryDataset(root: Optional[str] = None, 
                      transform: Optional[Callable] = None, 
                      pre_transform: Optional[Callable] = None, 
                      pre_filter: Optional[Callable] = None):
    """
        root: 根文件夹, 
            一个是raw_dir, 存放原始数据
            另一个是processed_dir, 存放处理后的数据
        transform: 此函数在每一次数据访问时被调用，所以它应该用于数据增广, 接受Data为参数
        pre_transform: 此函数在样本 Data对象保存到文件前调用，所以它最好用于只需要做一次的大量预计算。
        pre_filter: 函数可以在保存前手动过滤掉数据对象。该函数的一个用例是，过滤样本类别。
    """

创建一个自己的数据集类需要实现四个基本方法

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

import os
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 os.path.join(self.root, 'raw')

    @property
    def processed_dir(self):
        return os.path.join(self.root, 'processed')

    @property
    def raw_file_names(self):
        names = ['x', 'tx', 'allx', 'y', 'ty', 'ally', 'graph', 'test.index']
        return [f'ind.pubmed.{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(f'{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 f'{self.name}()'

dataset = PlanetoidPubMed('./input/Planetoid/PubMed')
print(dataset.num_classes)
# 3
print(dataset[0].num_nodes)
# 19717
print(dataset[0].num_edges)
# 88648
print(dataset[0].num_features)
# 500

构建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目录下是否存在pre_filter.pt文件，如果存在，意味着之前进行过样本过滤，则需加载该文件获取之前所用的样本过滤的方法，并检查它与当前pre_filter参数指定的方法是否相同，如果不相同则会报出一个警告，“The pre_filter argument differs from the one used in ……”。其中self.processed_dir为osp.join(self.root, 'processed')。
  • 接着检查是否存在处理好的数据：检查self.processed_dir目录下是否存在self.processed_paths方法返回的所有文件，如有文件不存在，意味着不存在已经处理好的样本的文件，如需执行以下的操作：
    • 调用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文件，其中分别存储当前使用的数据处理方法和样本过滤方法。

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

2.1 节点预测

• 定义模型

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

class GAT(torch.nn.Module):
    def __init__(self, num_features, hidden_channels_list, num_classes, seed=2021):
        super(GAT, self).__init__()
        torch.manual_seed(seed)
        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

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

• 准备数据

dataset = PlanetoidPubMed(root='./input/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)

• 训练，测试

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

for epoch in range(0, 300):
    loss = train()
    if epoch % 100 == 0:
        print(f'epcoh: {epoch+1:03d}, Loss: {loss:.4f}')

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

2.2 边预测

• 通过负采样, 采样一些不存在边的节点对作为负样本边，正负样本应平衡。
• train_test_split_edges(data, val_ratio=0.05, test_ratio=0.1)
  • data: 为PyG的Data对象
  • val_ratio: 为验证集比例
  • test_ratio: 为测试集比例
• 获取数据集

import os
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(root='./input/Cora', name="Cora", transform=T.NormalizeFeatures())
data = dataset[0]
print(data.edge_index.shape)
# torch.Size([2, 10556])

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)

• 构建模型
• 用于边预测的神经网络主要由 encode 和 decode 组成
  • encode: 与生成节点表征相同
  • decode: 边两端的节点的表征生成边为真的几率(odds)
• decode_all(self, z): 用于推断（inference）阶段，我们要对输入节点所有的节点对预测存在边的几率。

import torch
from torch_geometric.nn import GCNConv
from sklearn.metrics import roc_auc_score

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(Net, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, 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):
    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(model, data, 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(model, data):
    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(data):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    data = data.to(device)
    model = Net(dataset.num_features, 128, 64).to(device)
    optimizer = torch.optim.Adam(params=model.parameters(), lr=0.01)
    
    best_val_auc = test_auc = 0
    for epoch in range(0, 200):
        loss = train(model, data, optimizer)
        val_auc, tmp_test_auc = test(model, data)
        if val_auc > best_val_auc:
            best_val_auc = val_auc
            test_auc = tmp_test_auc
        if (epoch+1) % 10 == 0:
            print(f'epoch: {epoch+1}, loss: {loss:.4f}, val: {val_auc:.4f}, test: {test_auc:.4f}')
    z = model.encode(data.x, data.train_pos_edge_index)
    final_edge_index = model.decode_all(z)

main(data)

三、总结

3.1 InMemoryDataset 子类的运行流程与实现四个函数的规范

• 检查数据原始文件是否已下载
• 检查数据是否经过处理
  • 1.self.processed_dir目录下是否存在pre_transform.pt 文件
  • 2.检查之前的样本过滤的方法，检查self.processed_dir目录下是否存在pre_filter.pt 文件
  • 3.检查是否存在处理好的数据

3.2 构建模型

• 1.使用 GCNConv() 构建模型
• 2.使用 torch_geometric.nn.Sequential 构建模型

四、作业

• 实践问题一：对节点预测任务，尝试用PyG中的不同的网络层去代替GCNConv，以及不同的层数和不同的out_channels。

import torch
from torch.nn import Linear
import torch.nn.functional as F
from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, seed=2021):
        super(GAT, self).__init__()
        torch.manual_seed(seed)
        self.conv1 = GATConv(in_channels, hidden_channels)
        self.conv2 = GATConv(hidden_channels, out_channels)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

• 实践问题二：对边预测任务，尝试用用torch_geometric.nn.Sequential容器构造图神经网络。

from torch_geometric.nn import GCNConv, Sequential
from torch.nn import Linear, ReLU
import torch.nn.functional as F

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels_list, out_channels, seed=2021):
        super(Net, self).__init__()
        torch.manual_seed(seed)
        hns = [in_channels] + hidden_channels_list
        conv_list = []
        for idx in range(len(hidden_channels_list)):
            conv_list.append((GCNConv(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], out_channels)
        
    def encode(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 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()

五、参考资料

• DataWhale开源学习资料：https://github.com/datawhalechina/team-learning-nlp/tree/master/GNN