NNConv

class dgl.nn.mxnet.conv.NNConv(in_feats, out_feats, edge_func, aggregator_type, residual=False, bias=True)[source]

Bases: mxnet.gluon.block.Block

Graph Convolution layer from Neural Message Passing for Quantum Chemistry

\[h_{i}^{l+1} = h_{i}^{l} + \mathrm{aggregate}\left(\left\{ f_\Theta (e_{ij}) \cdot h_j^{l}, j\in \mathcal{N}(i) \right\}\right)\]

where \(e_{ij}\) is the edge feature, \(f_\Theta\) is a function with learnable parameters.

Parameters

  • in_feats (int) – Input feature size; i.e, the number of dimensions of \(h_j^{(l)}\). NN can be applied on homogeneous graph and unidirectional bipartite graph. If the layer is to be applied on a unidirectional bipartite graph, in_feats specifies the input feature size on both the source and destination nodes. If a scalar is given, the source and destination node feature size would take the same value.

  • out_feats (int) – Output feature size; i.e., the number of dimensions of \(h_i^{(l+1)}\).

  • edge_func (callable activation function/layer) – Maps each edge feature to a vector of shape (in_feats * out_feats) as weight to compute messages. Also is the \(f_\Theta\) in the formula.

  • aggregator_type (str) – Aggregator type to use (sum, mean or max).

  • residual (bool, optional) – If True, use residual connection. Default: False.

  • bias (bool, optional) – If True, adds a learnable bias to the output. Default: True.

Examples

>>> import dgl
>>> import numpy as np
>>> import mxnet as mx
>>> from mxnet import gluon
>>> from dgl.nn import NNConv
>>>
>>> # Case 1: Homogeneous graph
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = dgl.add_self_loop(g)
>>> feat = mx.nd.ones((6, 10))
>>> lin = gluon.nn.Dense(20)
>>> lin.initialize(ctx=mx.cpu(0))
>>> def edge_func(efeat):
>>>      return lin(efeat)
>>> efeat = mx.nd.ones((12, 5))
>>> conv = NNConv(10, 2, edge_func, 'mean')
>>> conv.initialize(ctx=mx.cpu(0))
>>> res = conv(g, feat, efeat)
>>> res
[[0.39946803 0.32098457]
[0.39946803 0.32098457]
[0.39946803 0.32098457]
[0.39946803 0.32098457]
[0.39946803 0.32098457]
[0.39946803 0.32098457]]
<NDArray 6x2 @cpu(0)>

>>> # Case 2: Unidirectional bipartite graph
>>> u = [0, 1, 0, 0, 1]
>>> v = [0, 1, 2, 3, 2]
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
>>> u_feat = mx.nd.random.randn(2, 10)
>>> v_feat = mx.nd.random.randn(4, 10)
>>> conv = NNConv(10, 2, edge_func, 'mean')
>>> conv.initialize(ctx=mx.cpu(0))
>>> efeat = mx.nd.ones((5, 5))
>>> res = conv(g, (u_feat, v_feat), efeat)
>>> res
[[ 0.24425688  0.3238042 ]
[-0.11651017 -0.01738572]
[ 0.06387337  0.15320925]
[ 0.24425688  0.3238042 ]]
<NDArray 4x2 @cpu(0)>
forward(graph, feat, efeat)[source]

Compute MPNN Graph Convolution layer.

Parameters

  • graph (DGLGraph) – The graph.

  • feat (mxnet.NDArray or pair of mxnet.NDArray) – The input feature of shape \((N, D_{in})\) where \(N\) is the number of nodes of the graph and \(D_{in}\) is the input feature size.

  • efeat (mxnet.NDArray) – The edge feature of shape \((N, *)\), should fit the input shape requirement of edge_nn.

Returns

The output feature of shape \((N, D_{out})\) where \(D_{out}\) is the output feature size.

Return type

mxnet.NDArray