参考链接：https://github.com/datawhalechina/team-learning-nlp/blob/master/GNN/Markdown%E7%89%88%E6%9C%AC/8-%E5%9F%BA%E4%BA%8E%E5%9B%BE%E7%A5%9E%E7%BB%8F%E7%BD%91%E7%BB%9C%E7%9A%84%E5%9B%BE%E8%A1%A8%E7%A4%BA%E5%AD%A6%E4%B9%A0.md
一、引言
图表征学习是在输入节点属性、边得到一个向量作为图的表征，基于图表征进一步做图的预测的一个学习过程。基于图同构网络（Graph Isomorphism Network, GIN）的图表征网络是当前最经典的图表征学习网络。
二、基于图同构网络（GIN）的图表征网络的实现
2.1 基于图同构网络的图表征学习的两个过程：
- 1.计算得到节点表征；
- 2.对图上各个节点的表征做图池化（Graph Pooling），或称为图读出（Graph Readout），得到图的表征（Graph Representation）。
2.2基于图同构网络的图表征模块（GINGraphRepr Module）
- 1.采用GINNodeEmbedding模块对图上每一个节点做节点嵌入（Node Embedding），得到节点表征；
- 2.对节点表征做图池化得到图的表征；
- 3.用一层线性变换对图表征转换为对图的预测。
代码如下：
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"):
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)
2.2基于图同构网络的节点嵌入模块（GINNodeEmbedding Module）
- 1.用AtomEncoder对其做嵌入得到第0层节点表征。
- 2.逐层计算节点表征，从第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
2.3 GINConv--图同构卷积层
图同构卷积层的数学定义如下：
PyG可以通过torch_geometric.nn.GINConv来使用PyG定义好的图同构卷积层，该实现不支持存在边属性的图。在这里自定义一个支持边属性的GINConv模块。
输入的边属性为类别型，需先将类别型边属性转换为边表征。定义的GINConv模块遵循“消息传递、消息聚合、消息更新”这一过程。
- 1.这一过程随着self.propagate()方法的调用开始执行，该函数接收edge_index, x, edge_attr此三个参数。接着message()方法被调用，在这里要传递的消息是源节点表征与边表征之和的relu()的输出。
- 2.在super(GINConv, self).init(aggr = "add")中定义了消息聚合方式为add，那么传入给任一个目标节点的所有消息被求和得到aggr_out，它还是目标节点的中间过程的信息。
- 3.接着执行消息更新过程，类GINConv继承了MessagePassing类，因此update()函数被调用。
我们希望对节点做消息更新中加入目标节点自身的消息，因此在update函数中只简单返回输入的aggr_out。然后在forward函数中，执行out = self.mlp((1 + self.eps) *x + self.propagate(edge_index, x=x, edge_attr=edge_embedding))实现消息的更新。
代码实现如下：
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
2.4AtomEncoder 与 BondEncoder
由于节点和边的属性都为离散值，它们属于不同的空间，无法直接将它们融合在一起。通过嵌入（Embedding），可以将节点属性和边属性分别映射到一个新的空间，在这个新的空间中，就可以对节点和边进行信息融合。在GINConv中，message()函数中的x_j + edge_attr 操作执行了节点信息和边信息的融合。
节点属性映射到一个新的空间是如何实现的?
- 1.full_atom_feature_dims 是一个链表list，存储了节点属性向量每一维可能取值的数量，即X[i] 可能的取值一共有full_atom_feature_dims[i]种情况，X为节点属性；
- 2.节点属性有多少维，那么就需要有多少个嵌入函数，通过调用torch.nn.Embedding(dim, emb_dim)可以实例化一个嵌入函数；
- 3.torch.nn.Embedding(dim, emb_dim)，第一个参数dim为被嵌入数据可能取值的数量，第一个参数emb_dim为要映射到的空间的维度。得到的嵌入函数接受一个大于0小于dim的数，输出一个维度为emb_dim的向量。嵌入函数也包含可训练参数，通过对神经网络的训练，嵌入函数的输出值能够表达不同输入值之间的相似性。
- 4.在forward()函数中，我们对不同属性值得到的不同嵌入向量进行相加操作，实现了将节点的的不同属性融合在一起。
代码实现如下：
import torch
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):
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)
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
if __name__ == '__main__':
from loader import GraphClassificationPygDataset
dataset = GraphClassificationPygDataset(name = 'tox21')
atom_enc = AtomEncoder(100)
bond_enc = BondEncoder(100)
print(atom_enc(dataset[0].x))
print(bond_enc(dataset[0].edge_attr))
三、理论分析
3.1动机（Motivation）
新的图神经网络的设计大多基于经验性的直觉、启发式的方法和实验性的试错。人们对图神经网络的特性和局限性了解甚少，对图神经网络的表征能力学习的正式分析也很有限。因此需要学习基于图神经网络的图表示学习。
3.2贡献与结论
1.（理论上）图神经网络在区分图结构方面最高能达到与WL Test一样的能力。
2.确定了邻接节点聚合方法和图池化方法应具备的条件，在这些条件下，所产生的图神经网络能达到与WL Test一样的能力。
3.分析过去流行的图神经网络变体（如GCN和GraphSAGE）无法区分一些结构的图。
4.开发了一个简单的图神经网络模型--图同构网络（Graph Isomorphism Network, GIN），并证明其分辨同构图的能力和表示图的能力与WL Test相当。
3.3背景：Weisfeiler-Lehman Test (WL Test)
图同构性测试
两个图是同构的，意思是两个图拥有一样的拓扑结构，即可以通过重新标记节点从一个图转换到另外一个图。Weisfeiler-Lehman 图的同构性测试算法，简称WL Test，是一种用于测试两个图是否同构的算法。
WL Test 的一维形式，类似于图神经网络中的邻接节点聚合。
1）WL Test 迭代地聚合节点及其邻接节点的标签；
2）将聚合的标签散列（hash）成新标签。该过程形式化为下方的公式， 在上方的公式中，表示节点的第次迭代的标签，第次迭代的标签为节点原始标签。
注意：
1.在迭代过程中，发现两个图之间的节点的标签不同时，就可以确定这两个图是非同构的。需要注意的是节点标签可能的取值只能是有限个数。
2.WL测试不能保证对所有图都有效，特别是对于具有高度对称性的图，如链式图、完全图、环图和星图，它会判断错误。
Weisfeiler-Lehman Graph Kernels 方法提出用WL子树核衡量图之间相似性。该方法使用WL Test不同迭代中的节点标签计数作为图的表征向量，它具有与WL Test相同的判别能力。直观地说，在WL Test的第次迭代中，一个节点的标签代表了以该节点为根的高度为的子树结构。
3.4图同构网络模型的构建
实现判断图同构性的图神经网络需要满足：两个节点自身标签一样且它们的邻接节点一样，这样图神经网络才可将这两个节点映射到相同的表征，即映射是单射性的。
可重复集合（Multisets）指的是元素可重复的集合，元素在集合中没有顺序关系。一个节点的所有邻接节点是一个可重复集合，一个节点可以有重复的邻接节点，邻接节点没有顺序关系。因此GIN模型中生成节点表征的方法遵循WL Test算法更新节点标签的过程。
在生成节点的表征后仍需要执行图池化（或称为图读出）操作得到图表征，最简单的图读出操作是做求和。由于每一层的节点表征都可能是重要的，因此在图同构网络中，不同层的节点表征在求和后被拼接，其数学定义如下， 采用拼接而不是相加的原因在于不同层节点的表征属于不同的特征空间。未做严格的证明，这样得到的图的表示与WL Subtree Kernel得到的图的表征是等价的。
四、结语
在此篇文章中，学习了基于图同构网络（GIN）的图表征网络，为了得到图表征首先需要做节点表征，然后做图读出。GIN中节点表征的计算遵循WL Test算法中节点标签的更新方法，因此它的上界是WL Test算法。在图读出中，对所有的节点表征（加权，如果用Attention的话）求和，这会造成节点分布信息的丢失。