基于图神经网络的图表征学习方法

1. 引言

图表征学习要求根据节点属性、边和边的属性（如果有的话）生成一个向量作为图的表征，基于图表征我们可以做图的预测。基于图同构网络（Graph Isomorphism Network, GIN）的图表征网络是当前最经典的图表征学习网络。

提出图同构网络的论文：How Powerful are Graph Neural Networks?

2. 基于图同构网络（GIN）的图表征网络的实现

基于图同构网络的图表征学习主要包含以下两个过程：

1. 首先计算得到节点表征；
2. 其次对图上各个节点的表征做图池化（Graph Pooling），或称为图读出（Graph Readout），得到图的表征（Graph Representation）。

AtomEncoder 与 BondEncoder

节点（原子）和边（化学键）的属性都为离散值，它们属于不同的空间，无法直接将它们融合在一起。通过嵌入（Embedding），我们可以将节点属性和边属性分别映射到一个新的空间，在这个新的空间中，我们就可以对节点和边进行信息融合。在GINConv中，message()函数中的x_j + edge_attr 操作执行了节点信息和边信息的融合。

通过下方的代码中的AtomEncoder类，来分析将节点属性映射到一个新的空间是如何实现的：

• full_atom_feature_dims 是一个链表list，存储了节点属性向量每一维可能取值的数量，即X[i]可能的取值一共有full_atom_feature_dims[i]种情况，X为节点属性；
• 节点属性有多少维，那么就需要有多少个嵌入函数，通过调用torch.nn.Embedding(dim, emb_dim)可以实例化一个嵌入函数；
• torch.nn.Embedding(dim, emb_dim)，第一个参数dim为被嵌入数据可能取值的数量，第二个参数emb_dim为要映射到的空间的维度。得到的嵌入函数接受一个大于0小于dim的数，输出一个维度为emb_dim的向量。嵌入函数也包含可训练参数，通过对神经网络的训练，嵌入函数的输出值能够表达不同输入值之间的相似性。
• 在forward()函数中，我们对不同属性值得到的不同嵌入向量进行了相加操作，实现了将节点的的不同属性融合在一起。

BondEncoder类与AtomEncoder类是类似的。

import torch
from torch import nn
from torch_geometric.nn import global_add_pool, global_mean_pool, global_max_pool, GlobalAttention, Set2Set

from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims 

full_atom_feature_dims = get_atom_feature_dims()
full_bond_feature_dims = get_bond_feature_dims()

class AtomEncoder(torch.nn.Module):
    """该类用于对原子属性做嵌入。
    记`N`为原子属性的维度，则原子属性表示为`[x1, x2, ..., xi, xN]`，其中任意的一维度`xi`都是类别型数据。full_atom_feature_dims[i]存储了原子属性`xi`的类别数量。
    该类将任意的原子属性`[x1, x2, ..., xi, xN]`转换为原子的嵌入`x_embedding`（维度为emb_dim）。
    """
    def __init__(self, emb_dim):
        super(AtomEncoder, self).__init__()
        
        self.atom_embedding_list = torch.nn.ModuleList()

        for i, dim in enumerate(full_atom_feature_dims):
            emb = torch.nn.Embedding(dim, emb_dim)  # 不同维度的属性用不同的Embedding方法
            torch.nn.init.xavier_uniform_(emb.weight.data)
            self.atom_embedding_list.append(emb)

    def forward(self, x):
        x_embedding = 0
        for i in range(x.shape[1]):
            x_embedding += self.atom_embedding_list[i](x[:,i])

        return x_embedding


class BondEncoder(torch.nn.Module):
    
    def __init__(self, emb_dim):
        super(BondEncoder, self).__init__()
        
        self.bond_embedding_list = torch.nn.ModuleList()

        for i, dim in enumerate(full_bond_feature_dims):
            emb = torch.nn.Embedding(dim, emb_dim)
            torch.nn.init.xavier_uniform_(emb.weight.data)
            self.bond_embedding_list.append(emb)

    def forward(self, edge_attr):
        bond_embedding = 0
        for i in range(edge_attr.shape[1]):
            bond_embedding += self.bond_embedding_list[i](edge_attr[:,i])

        return bond_embedding 

GINConv–图同构卷积层

图同构卷积层的数学定义如下： $${\mathbf{x}}_i^{\prime }=h_{\mathbf{\Theta }}$$