图网络消息传递框架¶

将卷积算子推广到不规则域 is typically expressed as a neighborhood aggregation or message passing scheme. 我们用 \(\mathbf{x}^{(k-1)}_i \in \mathbb{R}^F\) 表示第 \((k-1)\) 层 node \(i\) 的节点特征， \(\mathbf{e}_{j,i} \in \mathbb{R}^D\) 表示从node \(j\) 到 node \(i\) 的边特征，则消息传递图神经网络可以描述为

\[\mathbf{x}_i^{(k)} = \gamma^{(k)} \left( \mathbf{x}_i^{(k-1)}, \square_{j \in \mathcal{N}(i)} \, \phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)},\mathbf{e}_{j,i}\right) \right),\]

其中 \(\square\) 表示具有置换不变性的可微函数(例如 sum, mean or max, and \(\gamma\))， \(\phi\) 表示可微函数(例如多层感知机 MLPs).

“MessagePassing” 基类
GCN层实现
边卷积层的实现

“MessagePassing” 基类 ¶

PyTorch Geometric 提供了 torch_geometric.nn.MessagePassing 基类，来帮助创建消息传递图神经网络。因此我们只需要指定函数 \(\phi\) , i.e. message(), and \(\gamma\) , .i.e. update(), as well as the aggregation scheme to use, .i.e. aggr='add', aggr='mean' or aggr='max'.

This is done with the help of the following methods:

torch_geometric.nn.MessagePassing(aggr="add", flow="source_to_target"): Defines the aggregation scheme to use ("add", "mean" or "max") and the flow direction of message passing (either "source_to_target" or "target_to_source").
torch_geometric.nn.MessagePassing.propagate(edge_index, size=None, dim=0, **kwargs): The initial call to start propagating messages. Takes in the edge indices and all additional data which is needed to construct messages and to update node embeddings. Note that propagate() is not limited to exchange messages in symmetric adjacency matrices of shape [N, N] only, but can also exchange messages in general sparse assignment matrices, .e.g., bipartite graphs, of shape [N, M] by passing size=(N, M) as an additional argument. If set to None, the assignment matrix is assumed to be symmetric. For bipartite graphs with two independent sets of nodes and indices, and each set holding its own information, this split can be marked by passing the information as a tuple, e.g. x=(x_N, x_M). Furthermore, the dim attribute indicates along which axis to propagate.
torch_geometric.nn.MessagePassing.message(): Constructs messages to node \(i\) in analogy to \(\phi\) for each edge in \((j,i) \in \mathcal{E}\) if flow="source_to_target" and \((i,j) \in \mathcal{E}\) if flow="target_to_source". Can take any argument which was initially passed to propagate(). In addition, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.
torch_geometric.nn.MessagePassing.update(): Updates node embeddings in analogy to \(\gamma\) for each node \(i \in \mathcal{V}\). Takes in the output of aggregation as first argument and any argument which was initially passed to propagate().

Let us verify this by re-implementing two popular GNN variants, the GCN layer from Kipf and Welling and the EdgeConv layer from Wang et al..

GCN层实现 ¶

GCN layer 在数学上定义为

\[\mathbf{x}_i^{(k)} = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \frac{1}{\sqrt{\deg(i)} \cdot \sqrt{deg(j)}} \cdot \left( \mathbf{\Theta} \cdot \mathbf{x}_j^{(k-1)} \right),\]

where neighboring node features are first transformed by a weight matrix \(\mathbf{\Theta}\), normalized by their degree, and finally summed up. 该公式可以分为以下步骤：

Add self-loops to the adjacency matrix.
Linearly transform node feature matrix.
Compute normalization coefficients.
Normalize node features in \(\phi\).
Sum up neighboring node features ("add" aggregation).
Return new node embeddings in \(\gamma\).

通常在消息传递发生之前计算步骤1-3。使用 torch_geometric.nn.MessagePassing 基类可以轻松地执行步骤4-6。完整的层实现如下所示：

import torch
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(GCNConv, self).__init__(aggr='add')  # "Add" aggregation.
        self.lin = torch.nn.Linear(in_channels, out_channels)

    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(row, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-6: Start propagating messages.
        return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x,
                              norm=norm)

    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

    def update(self, aggr_out):
        # aggr_out has shape [N, out_channels]

        # Step 6: Return new node embeddings.
        return aggr_out

GCNConv inherits from torch_geometric.nn.MessagePassing with "add" propagation. All the logic of the layer takes place in forward(). Here, we first add self-loops to our edge indices using the torch_geometric.utils.add_self_loops() function (step 1), as well as linearly transform node features by calling the torch.nn.Linear instance (step 2).

We then proceed to call propagate(), which internally calls the message() and update() functions. As additional arguments for message propagation, we pass the node embeddings x.

In the message() function, we need to normalize the neighboring node features x_j. Here, x_j denotes a mapped tensor, which contains the neighboring node features of each edge. Node features can be automatically mapped by appending _i or _j to the variable name. In fact, any tensor can be mapped this way, as long as they have \(N\) entries in its first dimension.

The neighboring node features are normalized by computing node degrees \(\deg(i)\) for each node \(i\) and saving \(1/(\sqrt{\deg(i)} \cdot \sqrt{\deg(j)})\) in norm for each edge \((i,j) \in \mathcal{E}\).

In the update() function, we simply return the output of the aggregation.

这就是创建简单的消息传递层所需的全部。您可以将此层用作深度架构的组件。初始化和调用很简单：

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

边卷积层的实现 ¶

edge convolutional layer 用于处理图或点云，数学上定义为

\[\mathbf{x}_i^{(k)} = \max_{j \in \mathcal{N}(i)} h_{\mathbf{\Theta}} \left( \mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)} - \mathbf{x}_i^{(k-1)} \right),\]

其中 \(h_{\mathbf{\Theta}}\) 表示一个 MLP. 类似与 GCN layer, 我们也可以用 torch_geometric.nn.MessagePassing class 来实现 Edge Convolution, this time using the "max" aggregation:

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(EdgeConv, self).__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, size=(x.size(0), x.size(0)), 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)

    def update(self, aggr_out):
        # aggr_out has shape [N, out_channels]

        return aggr_out

Inside the message() function, we use self.mlp to transform both the target node features x_i and the relative source node features x_j - x_i for each edge \((j,i) \in \mathcal{E}\).

边卷积实际上是一种动态卷积，它使用特征空间中的最近邻居重新计算每一层的图。幸运的是，PyTorch Geometric comes with a GPU accelerated batch-wise k-NN graph generation method named torch_geometric.nn.knn_graph():

from torch_geometric.nn import knn_graph

class DynamicEdgeConv(EdgeConv):
    def __init__(self, in_channels, out_channels, k=6):
        super(DynamicEdgeConv, self).__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(DynamicEdgeConv, self).forward(x, edge_index)

Here, knn_graph() computes a nearest neighbor graph, which is further used to call the forward() method of EdgeConv.

这为我们提供了一个干净的接口，用于初始化和调用此层：

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

图网络消息传递框架¶

“MessagePassing” 基类¶

GCN层实现¶

边卷积层的实现¶

“MessagePassing” 基类 ¶

GCN层实现 ¶

边卷积层的实现 ¶