Task02 消息传递范式

pyG中的MessagePassing基类的简要定义

class MessagePassing(torch.nn.Module):
    def __init__(self, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2):
        super(MessagePassing, self).__init__()
		# 此处省略n行代码
        # Support for "fused" message passing.
        self.fuse = self.inspector.implements('message_and_aggregate')
		# 此处省略n行代码
    # 此处省略n行代码
    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
    	# 此处省略n行代码
        # Run "fused" message and aggregation (if applicable).
        if (isinstance(edge_index, SparseTensor) and self.fuse and not self.__explain__):
            coll_dict = self.__collect__(self.__fused_user_args__, edge_index, size, kwargs)

            msg_aggr_kwargs = self.inspector.distribute('message_and_aggregate', coll_dict)
            out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)

            update_kwargs = self.inspector.distribute('update', coll_dict)
            return self.update(out, **update_kwargs)
        # Otherwise, run both functions in separation.
        elif isinstance(edge_index, Tensor) or not self.fuse:
            coll_dict = self.__collect__(self.__user_args__, edge_index, size, kwargs)

            msg_kwargs = self.inspector.distribute('message', coll_dict)
            out = self.message(**msg_kwargs)
    		# 此处省略n行代码
            aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)
            out = self.aggregate(out, **aggr_kwargs)

            update_kwargs = self.inspector.distribute('update', coll_dict)
            return self.update(out, **update_kwargs)

可以看出，MessagePassing基类运行时，先执行__init__方法初始化，一个重要的flag变量self.fuse检查类定义是否实现了message_and_aggregate方法。MessagePassing执行时，消息传递的入口是propagate函数。进入propagate函数后，根据self.fuse这个指示变量以及传进来的数据是不是稀疏tensor(SparseTensor)决定走哪个分支：流程如下：

那么，我们使用pyG中的消息传递网络来设计自己的图神经网络模型时，只需要继承MessagePassing基类并且依次实现message，aggregate，和update函数就可以了。或者，把前两个合并，实现一个message_and_aggregate函数。

GCN的简单示例

Kipf的原论文中的公式为$X=RELU(A^{\ast }XW)$， 实验中，论文中的对称归一化$D^{-1/2}A{D}^{-1/2}$不如$D^{-1}A$直接归一化好使，在GCN原作者的pytorch实现中，也采用了后者的直接归一化。

GCN消息传递范式为：

$${\mathbf{x}}_i^{(k)}={\gamma }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\square }_{j\in \mathcal{N}(i)}{\varphi }^{(k)}\left({\mathbf{x}}_i^{(k-1)},{\mathbf{x}}_j^{(k-1)},{\mathbf{e}}_{j,i}\right)\right)$$

class GCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(GCNConv, self).__init__(aggr='add', flow='source_to_target')
        # "Add" aggregation (Step 5).
        # flow='source_to_target' 表示消息从源节点传播到目标节点
        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(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.
        return self.propagate(edge_index, 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