GNN入门之路06

最新推荐文章于 2023-01-18 17:14:20 发布

尼尔-冯-哈尔滨

最新推荐文章于 2023-01-18 17:14:20 发布

阅读量1.8k

点赞数

分类专栏：我的博客文章标签：人工智能 python 神经网络

本文链接：https://blog.csdn.net/m0_37671786/article/details/123580738

版权

我的博客专栏收录该内容

41 篇文章 2 订阅

订阅专栏

GNN入门之路06

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

在这一节的内容中我们将学习神经网络的图表征学习方法，我们要在已知节点属性、边的情况下生成一个向量作为图的表征。本文中的实现方法是最经典的图表征学习网络 图同构网络（Graph Isomorphism Network, GIN），我们将以它为例，通过该网络的实现、项目实践和理论分析，三个层面来学习基于图神经网络的图表征学习方法。
提出图同构网络的论文：How Powerful are Graph Neural Networks?

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

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

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

图同构网络的图表征模块（GINGraphRepr Module）

此模块的作用简单来说就是先将图的每一个节点表征计算出来，然后基于节点表征用池化的方法得到图的表征，最后用一层线性层变换来得到输出，代码如下：

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

class GINGraphRepr(nn.Module):

    def __init__(self, num_tasks=1, num_layers=5, emb_dim=300, residual=False, drop_ratio=0, JK="last", graph_pooling="sum"):
        """GIN Graph Pooling Module
        Args:
            num_tasks (int, optional): number of labels to be predicted. Defaults to 1 (控制了图表征的维度，dimension of graph representation).
            num_layers (int, optional): number of GINConv layers. Defaults to 5.
            emb_dim (int, optional): dimension of node embedding. Defaults to 300.
            residual (bool, optional): adding residual connection or not. Defaults to False.
            drop_ratio (float, optional): dropout rate. Defaults to 0.
            JK (str, optional): 可选的值为"last"和"sum"。选"last"，只取最后一层的结点的嵌入，选"sum"对各层的结点的嵌入求和。Defaults to "last".
            graph_pooling (str, optional): pooling method of node embedding. 可选的值为"sum"，"mean"，"max"，"attention"和"set2set"。 Defaults to "sum".

        Out:
            graph representation
        """
        super(GINGraphPooling, self).__init__()

        self.num_layers = num_layers
        self.drop_ratio = drop_ratio
        self.JK = JK
        self.emb_dim = emb_dim
        self.num_tasks = num_tasks

        if self.num_layers < 2:
            raise ValueError("Number of GNN layers must be greater than 1.")

        self.gnn_node = GINNodeEmbedding(num_layers, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual)

        # Pooling function to generate whole-graph embeddings
        if graph_pooling == "sum":
            self.pool = global_add_pool
        elif graph_pooling == "mean":
            self.pool = global_mean_pool
        elif graph_pooling == "max":
            self.pool = global_max_pool
        elif graph_pooling == "attention":
            self.pool = GlobalAttention(gate_nn=nn.Sequential(
                nn.Linear(emb_dim, emb_dim), nn.BatchNorm1d(emb_dim), nn.ReLU(), nn.Linear(emb_dim, 1)))
        elif graph_pooling == "set2set":
            self.pool = Set2Set(emb_dim, processing_steps=2)
        else:
            raise ValueError("Invalid graph pooling type.")

        if graph_pooling == "set2set":
            self.graph_pred_linear = nn.Linear(2*self.emb_dim, self.num_tasks)
        else:
            self.graph_pred_linear = nn.Linear(self.emb_dim, self.num_tasks)

    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)

可以看到可选的基于结点表征计算得到图表征的方法有：

“sum”：
- 对节点表征求和；
- 使用模块torch_geometric.nn.glob.global_add_pool。
“mean”：
- 对节点表征求平均；
- 使用模块torch_geometric.nn.glob.global_mean_pool。
“max”：取节点表征的最大值。
- 对一个batch中所有节点计算节点表征各个维度的最大值；
- 使用模块torch_geometric.nn.glob.global_max_pool。
“attention”：
- 基于Attention对节点表征加权求和；
- 使用模块 torch_geometric.nn.glob.GlobalAttention；
- 来自论文 “Gated Graph Sequence Neural Networks” 。
“set2set”:
1. 另一种基于Attention对节点表征加权求和的方法；
2. 使用模块 torch_geometric.nn.glob.Set2Set；
3. 来自论文 “Order Matters: Sequence to sequence for sets”。

PyG中集成的所有的图池化的方法可见于Global Pooling Layers。

接下来我们将学习节点嵌入的方法。

基于图同构网络的节点嵌入模块（GINNodeEmbedding Module）

此模块基于多层GINConv实现结点嵌入的计算。此处我们先忽略GINConv的实现。此模块得到的节点属性输入为类别型向量，我们首先用AtomEncoder对其做嵌入得到第0层节点表征（稍后我们再对AtomEncoder做分析）。然后我们逐层计算节点表征，从第1层开始到第num_layers层，每一层节点表征的计算都以上一层的节点表征h_list[layer]、边edge_index和边的属性edge_attr为输入。需要注意的是，GINConv的层数越多，此模块的感受野（receptive field）越大，结点i的表征最远能捕获到结点i的距离为num_layers的邻接节点的信息。

import torch
from mol_encoder import AtomEncoder
from gin_conv import GINConv
import torch.nn.functional as F

# GNN to generate node embedding
class GINNodeEmbedding(torch.nn.Module):
    """
    Output:
        node representations
    """

    def __init__(self, num_layers, emb_dim, drop_ratio=0.5, JK="last", residual=False):
        """GIN Node Embedding Module"""

        super(GINNodeEmbedding, self).__init__()
        self.num_layers = num_layers
        self.drop_ratio = drop_ratio
        self.JK = JK
        # add residual connection or not
        self.residual = residual

        if self.num_layers < 2:
            raise ValueError("Number of GNN layers must be greater than 1.")

        self.atom_encoder = AtomEncoder(emb_dim)

        # List of GNNs
        self.convs = torch.nn.ModuleList()
        self.batch_norms = torch.nn.ModuleList()

        for layer in range(num_layers):
            self.convs.append(GINConv(emb_dim))
            self.batch_norms.append(torch.nn.BatchNorm1d(emb_dim))

    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

接下来我们来学习图同构网络的关键组件GINConv。

`GINConv`–图同构卷积层

图同构卷积层的数学定义如下：
$\mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot\mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \right)$
PyG中已经实现了此模块，我们可以通过torch_geometric.nn.GINConv来使用PyG定义好的图同构卷积层，然而该实现不支持存在边属性的图。在这里我们自己自定义一个支持边属性的GINConv模块。

import torch
from torch import nn
from torch_geometric.nn import MessagePassing
import torch.nn.functional as F
from ogb.graphproppred.mol_encoder import BondEncoder


### GIN convolution along the graph structure
class GINConv(MessagePassing):
    def __init__(self, emb_dim):
        '''
            emb_dim (int): node embedding dimensionality
        '''
        super(GINConv, self).__init__(aggr = "add")

        self.mlp = nn.Sequential(nn.Linear(emb_dim, emb_dim), nn.BatchNorm1d(emb_dim), nn.ReLU(), nn.Linear(emb_dim, emb_dim))
        self.eps = nn.Parameter(torch.Tensor([0]))
        self.bond_encoder = BondEncoder(emb_dim = emb_dim)

    def forward(self, x, edge_index, edge_attr):
        edge_embedding = self.bond_encoder(edge_attr) # 先将类别型边属性转换为边表征
        out = self.mlp((1 + self.eps) *x + self.propagate(edge_index, x=x, edge_attr=edge_embedding))
        return out

    def message(self, x_j, edge_attr):
        return F.relu(x_j + edge_attr)
        
    def update(self, aggr_out):
        return aggr_out

理论分析

动机

论文原话：“The design of new GNNs is mostly based on empirical intuition, heuristics, and experimental trial-and-error. There is little understanding of the properties and limitations of GNNs, and formal analysis of GNNs’ representational capacity is limited.” 新的图神经网络的设计大多基于经验性的直觉、启发式的方法和实验性的试错。人们对图神经网络的特性和局限性了解甚少，对图神经网络的表征能力学习的正式分析也很有限。