Datawhale 6月学习——图神经网络：图预测任务实践

前情回顾

1. 图神经网络：图数据表示及应用
2. 图神经网络：消息传递图神经网络
3. 图神经网络：基于GNN的节点表征学习
4. 图神经网络：基于GNN的节点预测任务及边预测任务
5. 图神经网络：超大图上的节点表征学习
6. 图神经网络：基于图神经网络的图表征学习

1 图预测任务描述

1.1 任务简述

本次任务所涉及的图预测任务，是一个回归任务。其目的是，学习图结构的一个合理的图表征方式，学习目标的检测方式是与benchmark数据集上打好的数值标签进行对比。
这个数值标签可以是药物疗效评价等。

具体一点描述的话，就是对某一个已有标签(y)的数据集，我们要学习一种图表征的模型，使得这个图表征模型的计算结果与标签(y)最接近，即回归任务。

1.2 数据集

本次所用数据集，是OGB LSC中Graph level的数据集PCQM4M-LSC，是一个量子化学数据集，任务是预测给定分子的重要分子特性，即 HOMO-LUMO 间隙（图形回归）。
这个数据集的详细介绍可以看OGB官网。

这个数据的下载大小为58MB，但是是以smile字符串形式存储，将它们处理成图形对象后，最终的文件大小将在 8GB 左右，因此需要采用合适的数据读取形式来减少内存负担。

2 任务实现

本次任务分为两个部分，一个部分是按需获取数据集类的创建，第二个部分是图表示学习的实现。

2.1 数据集创建及获取

2.1.1 按需获取数据集类的创建

在前面的学习中我们只接触了数据可全部储存于内存的数据集，这些数据集对应的数据集类在创建对象时就将所有数据都加载到内存。然而在一些应用场景中，数据集规模超级大，我们很难有足够大的内存完全存下所有数据。因此需要一个按需加载样本到内存的数据集类。在此上半节内容中，我们将学习为一个包含上千万个图样本的数据集构建一个数据集类。

共有三种方法可以实现这样的数据集类的创建

1. 通过继承torch_geometric.data.Dataset基类来自定义一个按需加载样本到内存的数据集类

import os.path as osp

import torch
from torch_geometric.data import Dataset, download_url

class MyOwnDataset(Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyOwnDataset, self).__init__(root, transform, pre_transform)

    @property
    def raw_file_names(self):
        return ['some_file_1', 'some_file_2', ...]

    @property
    def processed_file_names(self):
        return ['data_1.pt', 'data_2.pt', ...]

    def download(self):
        # Download to `self.raw_dir`.
        path = download_url(url, self.raw_dir)
        ...

    def process(self):
        i = 0
        for raw_path in self.raw_paths:
            # Read data from `raw_path`.
            data = Data(...)

            if self.pre_filter is not None and not self.pre_filter(data):
                continue

            if self.pre_transform is not None:
                data = self.pre_transform(data)

            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data

2. 直接生成一个Dataloader对象

from torch_geometric.data import Data, DataLoader

data_list = [Data(...), ..., Data(...)]
loader = DataLoader(data_list, batch_size=32)

3. 将一个列表的Data对象组成一个batch

from torch_geometric.data import Data, Batch

data_list = [Data(...), ..., Data(...)]
loader = Batch.from_data_list(data_list, batch_size=32)
#上述代码运行失败，暂未解决

接下来，也可对图样本封装成批（BATCHING）与DataLoader类，有几种方式

1. 合并小图成大图
2. 小图的属性增值与拼接

2.1.2 创建预测任务所需数据集

我们定义的数据集类如下：

import os
import os.path as osp

import pandas as pd
import torch
from ogb.utils.mol import smiles2graph
from ogb.utils.torch_util import replace_numpy_with_torchtensor
from ogb.utils.url import download_url, extract_zip
from rdkit import RDLogger
from torch_geometric.data import Data, Dataset
import shutil

RDLogger.DisableLog('rdApp.*')

class MyPCQM4MDataset(Dataset):

    def __init__(self, root):
        self.url = 'https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m_kddcup2021.zip'
        super(MyPCQM4MDataset, self).__init__(root)

        filepath = osp.join(root, 'raw/data.csv.gz')
        data_df = pd.read_csv(filepath)
        self.smiles_list = data_df['smiles']
        self.homolumogap_list = data_df['homolumogap']

    @property
    def raw_file_names(self):
        return 'data.csv.gz'

    def download(self):
        path = download_url(self.url, self.root)
        extract_zip(path, self.root)
        os.unlink(path)
        shutil.move(osp.join(self.root, 'pcqm4m_kddcup2021/raw/data.csv.gz'), osp.join(self.root, 'raw/data.csv.gz'))

    def len(self):
        return len(self.smiles_list)

    def get(self, idx):
        smiles, homolumogap = self.smiles_list[idx], self.homolumogap_list[idx]
        graph = smiles2graph(smiles)
        assert(len(graph['edge_feat']) == graph['edge_index'].shape[1])
        assert(len(graph['node_feat']) == graph['num_nodes'])

        x = torch.from_numpy(graph['node_feat']).to(torch.int64)
        edge_index = torch.from_numpy(graph['edge_index']).to(torch.int64)
        edge_attr = torch.from_numpy(graph['edge_feat']).to(torch.int64)
        y = torch.Tensor([homolumogap])
        num_nodes = int(graph['num_nodes'])
        data = Data(x, edge_index, edge_attr, y, num_nodes=num_nodes)
        return data

    # 获取数据集划分
    def get_idx_split(self):
        split_dict = replace_numpy_with_torchtensor(torch.load(osp.join(self.root, 'pcqm4m_kddcup2021/split_dict.pt')))
        return split_dict

if __name__ == "__main__":
    dataset = MyPCQM4MDataset('dataset2')
    from torch_geometric.data import DataLoader
    from tqdm import tqdm
    dataloader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=4)
    for batch in tqdm(dataloader):
        pass

以上代码依赖于ogb包，通过pip install ogb命令可安装此包。ogb文档可见于Get Started | Open Graph Benchmark (stanford.edu)。

在生成一个该数据集类的对象时，程序首先会检查指定的文件夹下是否存在data.csv.gz文件，如果不在，则会执行download方法，这一过程是在运行super类的__init__方法中发生的。然后程序继续执行__init__方法的剩余部分，读取data.csv.gz文件，获取存储图信息的smiles格式的字符串，以及回归预测的目标homolumogap。我们将由smiles格式的字符串转成图的过程在get()方法中实现，这样我们在生成一个DataLoader变量时，通过指定num_workers可以实现并行执行生成多个图。

在使用colab进行ogb包的调用时，提示未找到RDKit，这是一个化学信息学的开源工具包，需要借助conda环境进行配置。
由于colab环境不预先配有conda，需要进行安装，此处参考StackOverflow问答Installing RDKit in Google Colab，使用以下命令进行安装配置。

!pip install -q condacolab
import condacolab
condacolab.install()

!mamba install -c conda-forge rdkit

2.2 图表征模型的建立

此部分的代码详见上一任务 图神经网络：基于图神经网络的图表征学习，此处再进行一个细化的理解。重点在于节点嵌入的实现GINNodeEmbedding及图池化的实现GINPoolingRepr

2.2.1 节点嵌入的实现

class GINNodeEmbedding(torch.nn.Module):
...
    def forward(self, batched_data):
        x, edge_index, edge_attr = batched_data.x, batched_data.edge_index, batched_data.edge_attr

        # computing input node embedding
        h_list = [self.atom_encoder(x)]  # 先将类别型原子属性转化为原子表征
        for layer in range(self.num_layers):
            h = self.convs[layer](h_list[layer], edge_index, edge_attr)
            h = self.batch_norms[layer](h)
            if layer == self.num_layers - 1:
                # remove relu for the last layer
                h = F.dropout(h, self.drop_ratio, training=self.training)
            else:
                h = F.dropout(F.relu(h), self.drop_ratio, training=self.training)

            if self.residual:
                h += h_list[layer]

            h_list.append(h)

        # Different implementations of Jk-concat
        if self.JK == "last":
            node_representation = h_list[-1]
        elif self.JK == "sum":
            node_representation = 0
            for layer in range(self.num_layers + 1):
                node_representation += h_list[layer]

        return node_representation

实现主要功能的forward函数，是由atom_encoder层，若干GCNConv及一个dropout层构成的，最后将实现一个结果的聚合（通过最后一层的形式last，或者求和的形式进行结果聚合sum）。

2.2.2 图池化的实现

class GINGraphRepr(nn.Module):
...
    def forward(self, batched_data):
        h_node = self.gnn_node(batched_data)

        h_graph = self.pool(h_node, batched_data.batch)
        output = self.graph_pred_linear(h_graph)

        if self.training:
            return output
        else:
            # At inference time, relu is applied to output to ensure positivity
            # 因为预测目标的取值范围就在 (0, 50] 内
            return torch.clamp(output, min=0, max=50)

对节点嵌入结果先进行池化，再进行线性变换到想要的输出个数h_graph。

2.3 预测任务的实现

2.3.1 参数设定

这边使用了argparse进行参数设定，具体可以查看官方文档。

import argparse
def parse_args():

    parser = argparse.ArgumentParser(description='Graph data miming with GNN')
    parser.add_argument('--task_name', type=str, default='GINGraphPooling',
                        help='task name')
    parser.add_argument('--device', type=int, default=0,
                        help='which gpu to use if any (default: 0)')
    parser.add_argument('--num_layers', type=int, default=5,
                        help='number of GNN message passing layers (default: 5)')
    parser.add_argument('--graph_pooling', type=str, default='sum',
                        help='graph pooling strategy mean or sum (default: sum)')
    parser.add_argument('--emb_dim', type=int, default=256,
                        help='dimensionality of hidden units in GNNs (default: 256)')
    parser.add_argument('--drop_ratio', type=float, default=0.,
                        help='dropout ratio (default: 0.)')
    parser.add_argument('--save_test', action='store_true')
    parser.add_argument('--batch_size', type=int, default=512,
                        help='input batch size for training (default: 512)')
    parser.add_argument('--epochs', type=int, default=100,
                        help='number of epochs to train (default: 100)')
    parser.add_argument('--weight_decay', type=float, default=0.00001,
                        help='weight decay')
    parser.add_argument('--early_stop', type=int, default=10,
                        help='early stop (default: 10)')
    parser.add_argument('--num_workers', type=int, default=0,#4,
                        help='number of workers (default: 4)')
    parser.add_argument('--dataset_root', type=str, default="dataset",
                        help='dataset root')
    args = parser.parse_args()

    return args

由于使用了parse_args，直接在ipython上运行（如colab）会导致报错，参考StackOverflow问答SystemExit: 2 error when calling parse_args()，在代码中添加下述代码，可以正常执行

import sys
sys.argv=['']
del sys

2.3.2 训练函数，求解函数及测试函数的定义

def train(model, device, loader, optimizer, criterion_fn):
    model.train()
    loss_accum = 0

    for step, batch in enumerate(tqdm(loader)):
        batch = batch.to(device)
        pred = model(batch).view(-1,)
        optimizer.zero_grad()
        loss = criterion_fn(pred, batch.y)
        loss.backward()
        optimizer.step()
        loss_accum += loss.detach().cpu().item()

    return loss_accum / (step + 1)


def eval(model, device, loader, evaluator):
    model.eval()
    y_true = []
    y_pred = []

    with torch.no_grad():
        for _, batch in enumerate(tqdm(loader)):
            batch = batch.to(device)
            pred = model(batch).view(-1,)
            y_true.append(batch.y.view(pred.shape).detach().cpu())
            y_pred.append(pred.detach().cpu())

    y_true = torch.cat(y_true, dim=0)
    y_pred = torch.cat(y_pred, dim=0)
    input_dict = {"y_true": y_true, "y_pred": y_pred}
    return evaluator.eval(input_dict)["mae"]


def test(model, device, loader):
    model.eval()
    y_pred = []

    with torch.no_grad():
        for _, batch in enumerate(loader):
            batch = batch.to(device)
            pred = model(batch).view(-1,)
            y_pred.append(pred.detach().cpu())

    y_pred = torch.cat(y_pred, dim=0)
    return y_pred

此处没有什么特别的逻辑，故不展开讨论。

2.3.3 训练

首先定义结果存储位置，计算设备等信息

def prepartion(args):
    save_dir = os.path.join('saves', args.task_name)
    if os.path.exists(save_dir):
        for idx in range(1000):
            if not os.path.exists(save_dir + '=' + str(idx)):
                save_dir = save_dir + '=' + str(idx)
                break

    args.save_dir = save_dir
    os.makedirs(args.save_dir, exist_ok=True)
    args.device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
    args.output_file = open(os.path.join(args.save_dir, 'output'), 'a')
    print(args, file=args.output_file, flush=True)

定义主函数

def main(args):
    prepartion(args)
    nn_params = {
        'num_layers': args.num_layers,
        'emb_dim': args.emb_dim,
        'drop_ratio': args.drop_ratio,
        'graph_pooling': args.graph_pooling
    }

    # automatic dataloading and splitting
    dataset = MyPCQM4MDataset(root=args.dataset_root)
    split_idx = dataset.get_idx_split()
    train_data = dataset[split_idx['train']]
    valid_data = dataset[split_idx['valid']]
    test_data = dataset[split_idx['test']]
    train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
    valid_loader = DataLoader(valid_data, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
    test_loader = DataLoader(test_data, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)

    # automatic evaluator. takes dataset name as input
    evaluator = PCQM4MEvaluator()
    criterion_fn = torch.nn.MSELoss()

    device = args.device

    model = GINGraphPooling(**nn_params).to(device)

    num_params = sum(p.numel() for p in model.parameters())
    print(f'#Params: {num_params}', file=args.output_file, flush=True)
    print(model, file=args.output_file, flush=True)

    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=args.weight_decay)
    scheduler = StepLR(optimizer, step_size=30, gamma=0.25)

    writer = SummaryWriter(log_dir=args.save_dir)
    not_improved = 0
    best_valid_mae = 9999

    for epoch in range(1, args.epochs + 1):
        print("=====Epoch {}".format(epoch), file=args.output_file, flush=True)
        print('Training...', file=args.output_file, flush=True)
        train_mae = train(model, device, train_loader, optimizer, criterion_fn)

        print('Evaluating...', file=args.output_file, flush=True)
        valid_mae = eval(model, device, valid_loader, evaluator)

        print({'Train': train_mae, 'Validation': valid_mae}, file=args.output_file, flush=True)

        writer.add_scalar('valid/mae', valid_mae, epoch)
        writer.add_scalar('train/mae', train_mae, epoch)

        if valid_mae < best_valid_mae:
            best_valid_mae = valid_mae
            if args.save_test:
                print('Saving checkpoint...', file=args.output_file, flush=True)
                checkpoint = {
                    'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict(), 'best_val_mae': best_valid_mae, 'num_params': num_params
                }
                torch.save(checkpoint, os.path.join(args.save_dir, 'checkpoint.pt'))
                print('Predicting on test data...', file=args.output_file, flush=True)
                y_pred = test(model, device, test_loader)
                print('Saving test submission file...', file=args.output_file, flush=True)
                evaluator.save_test_submission({'y_pred': y_pred}, args.save_dir)

            not_improved = 0
        else:
            not_improved += 1
            if not_improved == args.early_stop:
                print(f"Have not improved for {not_improved} epoches.", file=args.output_file, flush=True)
                break

        scheduler.step()
        print(f'Best validation MAE so far: {best_valid_mae}', file=args.output_file, flush=True)

    writer.close()
    args.output_file.close()

上述函数分为几个主要的部分

1. 数据获取，分片
2. 生成模型对象，验证器对象，优化函数等
3. 定义tensorboard写入对象SummaryWriter()
4. 进行训练，并对每一步都存储训练结果和验证结果，写入到SummaryWriter()中
   其中，定义了早停逻辑，当验证集上的最佳mae，超过args.early_stop次数不再变化时，会触发早停，训练结束。

2.4 训练及结果

2.4.1 训练设备及速度

首先在colab GPU 高RAM模式下进行了训练，由于涉及到的读写内容很多，速度很慢

等不下去了，更换本地电脑设备进行计算。有趣的是，虽然是windows系统设备，但成功开启了num_workers=4的计算大门。
计算耗时4个多小时。

2.4.2 训练结果

将训练结果在tensorboard中打开。
通过anaconda prompt，输入下面代码。

tensorboard --logdir=......\\gin_regression\\saves

会得到

TensorFlow installation not found - running with reduced feature set.
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.5.0 at http://localhost:6006/ (Press CTRL+C to quit)

复制网址到浏览器，可以看到结果。
在训练集上的mae变化：

在验证集上的mae变化：

由于训练前忘了将model存出，而训练耗时很长，此处没有来得及计算测试集上的mae情况。

参考阅读

1. Datawhale组队学习
2. 图神经网络在分子性质预测任务中的应用