【图神经网络】Pytorch图神经网络库——PyG创建消息传递网络

将卷积算子推广到不规则域通常表示为邻域聚合或消息传递方案。在第 $(k-1)$层中节点 $i$的节点特征用 ${\mathrm{x}}_i^{(k-1)}\in {\mathbb{R}}^F$表示，从节点 $j$到节点 $i$的边特征用 ${\mathrm{e}}_{j,i}\in {\mathbb{R}}^D$表示，消息传递图神经网络可以用以下公式描述：
$${\mathrm{x}}_i^{(k)}={\gamma }^{(k)}({\mathrm{x}}_i^{(k-1)},{\square }_{j\in \mathcal{N}(i)}{\varphi }^{(k)}({\mathrm{x}}_i^{(k-1)},{\mathrm{x}}_j^{(k-1)},{\mathrm{e}}_{j,i}))$$
其中 $\square$表示一个可微、具有对称性的函数，比如sum，max，mean， $\gamma$和 $\varphi$表示不同的函数，比如MLPs。

消息传递基类：MessagePassing

PyG提供了消息传递基类，用于创建GNN自动化的消息传递机制。用户只需要定义函数$\gamma$和$\varphi$，分别表示message()和update()。聚合操作有aggr="add", aggr="mean" or aggr="max"等。

下面是一些相关方法的简介：
MessagePassing(aggr="add", flow="source_to_target", node_dim=-2)：定义了一个聚集机制，三个参数分别为：聚集方式，消息传递方向以及沿哪个维度进行传播。
MessagePassing.propagate(edge_index, size=None, **kwargs)：首次调用开始传播消息。获取边索引edge_index和所有用于构造消息和更新节点嵌入的附加数据。这个函数不但可以用于方阵，而且也可以用于二分图等非方阵图，但是需要传递size参数表明矩阵形状size=(N, M)。
MessagePassing.message(...)：构造消息到节点$i$，但是根据传播方向有两种情况，如果边方向是$(j,i)$ 且flow="source_to_target"，即边是$j$指向$i$，而且消息流向是源节点到目的节点，或者相反。通常将中心节点表示为$i$，邻居节点表示为$j$。
MessagePassing.update(aggr_out, ...)：更新每个节点$i$的嵌入向量，接受聚合的输出作为第一个参数以及最初传递给propagate()的任何参数。

总之，PyG要么是直接调用nn里面的层，或者自己实现网络层。调用nn里面的层在上次简介介绍过了，下面就看一下如何使用MessagePassing这个基类继承实现GCN和EdgeConv层。

GCN层的实现

GCN层的数学定义如下：
$${\mathrm{x}}_i^{(k)}=\sum _{j\in \mathrm{N}(i)\cup \{i\}}\frac{1}{\sqrt{deg(i)}\cdot \sqrt{deg(j)}}\cdot ({\Theta }^T\cdot {\mathrm{x}}_j^{(k-1)})+\mathrm{b}$$
其中邻居节点的特征首先通过权重矩阵$\Theta$进行变换，然后使用它们的度进行标准化（normalized），最终加和（summed up）。将其步骤写为如下几步：

1. 向邻接矩阵(adjacency matrix)添加自循环(self-loops)；
2. 线性变换节点特征矩阵；
3. 计算归一化系数；
4. 在$\varphi$中Normalize节点特性
5. 对相邻节点特征进行归纳（add聚合）
6. 应用一个最终的偏差向量

步骤1-3通常是在消息传递之前计算的。使用MessagePassing基类可以很容易地处理步骤4-5。完整的GCN层实现如下：

import torch
from torch.nn import Linear, Parameter
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "Add" aggregation (Step 5).
        self.lin = Linear(in_channels, out_channels, bias=False)
        self.bias = Parameter(torch.Tensor(out_channels))

        self.reset_parameters()

    def reset_parameters(self):
        self.lin.reset_parameters()
        self.bias.data.zero_()

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: Linearly transform node feature matrix.
        x = self.lin(x)

        # Step 3: Compute normalization.
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-5: Start propagating messages.
        out = self.propagate(edge_index, x=x, norm=norm)

        # Step 6: Apply a final bias vector.
        out += self.bias

        return out

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]

        # Step 4: Normalize node features.
        return norm.view(-1, 1) * x_j

GCNConv通过“add”传播继承自MessagePassing。该层的所有逻辑都发生在其forward（）方法中。这里，我们首先使用torch_geometric.utils.add_self_loops（）函数将自循环添加到边索引中（步骤1），以及通过调用torch.nn.Linear实例来线性变换节点特征（步骤2）。
对于每一个节点$i$，归一化系数由节点度$deg(i)$得出，对于每个边$(j,i)\in \mathcal{E}$，转换到了$1/(\sqrt{deg(i)}\cdot \sqrt{deg(j)})$,结果保存在形为[ num _ edge，](步骤3)的张量norm中。

然后，我们调用propagate()，该函数在内部调用message()，aggregate()和update()。我们将节点嵌入$x$和标准化系数norm作为消息传播的其他参数(additional arguments)。

在message()函数中，我们需要通过norm规范相邻节点的特征x_j。这里，x_j表示lifted张量，它包含每个边的源节点特征，即每个节点的邻居。通过将_i或_j附加到变量名称，可以自动提升节点特征。事实上，任何张量都可以这样转换，只要它们包含源节点或目标节点特征。

这就是创建一个简单的消息传递层所需的全部内容。可以将此层用作深层体系结构的构建块。初始化和调用它非常简单:

conv = GCNConv(16, 32)
x = conv(x, edge_index)

实现Edge Convolution

Edge Convolution通过下式进行图或点云的处理：
$${\mathrm{x}}_i^{(k)}=\underset{j\in \mathrm{N}(i)}{\max }{h}_{\Theta }({\mathrm{x}}_i^{(k-1)},{\mathrm{x}}_j^{(k-1)}-{\mathrm{x}}_i^{(k-1)})$$
其中$h_{\Theta }$表示一个MLP，类比GCN，我们使用MessagePassing 来实现这一层，同时使用max聚合。实现代码如下：

import torch
from torch.nn import Sequential as Seq, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='max') #  "Max" aggregation.
        self.mlp = Seq(Linear(2 * in_channels, out_channels),
                       ReLU(),
                       Linear(out_channels, out_channels))

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        return self.propagate(edge_index, x=x)

    def message(self, x_i, x_j):
        # x_i has shape [E, in_channels]
        # x_j has shape [E, in_channels]

        tmp = torch.cat([x_i, x_j - x_i], dim=1)  # tmp has shape [E, 2 * in_channels]
        return self.mlp(tmp)

在message()函数内部，对于每个边$(j,i)\in \mathcal{E}$，我们使用self.mlp将节点特征$x_i$和相对源节点特征$x_j-x_i$进行转换。

根据定义式，就是首先使用MLP处理输入，在使用Max的aggregation操作。其中，MLP需要自己定义，然后aggregation操作只需要在初始化父类时传入参数aggr='max'即可，非常的方便。

而对于边的卷积操作，实际上是动态的卷积，在每一层使用最近邻居在特征空间进行重计算。幸运的是，PyG有一个使用GPU加速的K-NN图产生的方法torch_geometric.nn.pool.knn_graph()。

from torch_geometric.nn import knn_graph

class DynamicEdgeConv(EdgeConv):
    def __init__(self, in_channels, out_channels, k=6):
        super().__init__(in_channels, out_channels)
        self.k = k

    def forward(self, x, batch=None):
        edge_index = knn_graph(x, self.k, batch, loop=False, flow=self.flow)
        return super().forward(x, edge_index)

knn_graph方法计算最近邻图，然后进一步调用EdgeConv的forward函数进行传播。至此，就可以为DynamicEdgeConv留下一个干净的调用接口。

conv = DynamicEdgeConv(3, 128, k=6)
x = conv(x, batch)

内容来源：

[1] CREATING MESSAGE PASSING NETWORKS