EGNNConv

class dgl.nn.pytorch.conv.EGNNConv(in_size, hidden_size, out_size, edge_feat_size=0)[source]

基类： Module

\[ \begin{align}\begin{aligned}m_{ij}=\phi_e(h_i^l, h_j^l, ||x_i^l-x_j^l||^2, a_{ij})\\x_i^{l+1} = x_i^l + C\sum_{j\in\mathcal{N}(i)}(x_i^l-x_j^l)\phi_x(m_{ij})\\m_i = \sum_{j\in\mathcal{N}(i)} m_{ij}\\h_i^{l+1} = \phi_h(h_i^l, m_i)\end{aligned}\end{align} \]

其中 \(h_i\)、\(x_i\) 和 \(a_{ij}\) 分别是节点特征、坐标特征和边特征。\(\phi_e\)、\(\phi_h\) 和 \(\phi_x\) 是两层 MLP。\(C\) 是用于归一化的常数，计算方式为 \(1/|\mathcal{N}(i)|\)。

参数:

in_size (int) – 输入特征大小；即 \(h_i^l\) 的大小。
hidden_size (int) – 隐藏特征大小；即 \(\phi_e, \phi_x, \phi_h\) 中两层 MLP 里隐藏层的大小。
out_size (int) – 输出特征大小；即 \(h_i^{l+1}\) 的大小。
edge_feat_size (int, optional) – 边特征大小；即 \(a_{ij}\) 的大小。默认为 0。

示例

>>> import dgl
>>> import torch as th
>>> from dgl.nn import EGNNConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> node_feat, coord_feat, edge_feat = th.ones(6, 10), th.ones(6, 3), th.ones(6, 2)
>>> conv = EGNNConv(10, 10, 10, 2)
>>> h, x = conv(g, node_feat, coord_feat, edge_feat)

forward(graph, node_feat, coord_feat, edge_feat=None)[source]

描述

计算 EGNN 层。

参数 graph:

图。

类型 graph:

DGLGraph

参数 node_feat:

输入特征，形状为 \((N, h_n)\)。\(N\) 是节点数，\(h_n\) 必须与 in_size 相同。

类型 node_feat:

torch.Tensor

参数 coord_feat:

坐标特征，形状为 \((N, h_x)\)。\(N\) 是节点数，\(h_x\) 可以是任意正整数。

类型 coord_feat:

torch.Tensor

参数 edge_feat:

边特征，形状为 \((M, h_e)\)。\(M\) 是边数，\(h_e\) 必须与 edge_feat_size 相同。

类型 edge_feat:

torch.Tensor, 可选

返回值:

node_feat_out (torch.Tensor) – 输出节点特征，形状为 \((N, h_n')\)，其中 \(h_n'\) 与 out_size 相同。
coord_feat_out (torch.Tensor) – 输出坐标特征，形状为 \((N, h_x)\)，其中 \(h_x\) 与输入的坐标特征维度相同。