mindspore_gl.nn.NNConv

View Source On Gitee
class mindspore_gl.nn.NNConv(in_feat_size: int, out_feat_size: int, edge_embed, aggregator_type: str = 'sum', residual=False, bias=True)[source]

Graph convolutional layer. From the paper 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 \(f_\Theta\) is a function with learnable parameters.

Parameters

  • in_feat_size (int) – Input node feature size.

  • out_feat_size (int) – Output node feature size.

  • edge_embed (mindspore.nn.Cell) – Edge embedding function Cell.

  • aggregator_type (str, optional) – Type of aggregator, should be 'sum'. Default: 'sum'.

  • residual (bool, optional) – Whether use residual. Default: False.

  • bias (bool, optional) – Whether use bias. Default: True.

Inputs:

  • x (Tensor) - The input node features. The shape is \((N,D\_in)\) where \(N\) is the number of nodes and \(D\_in\) could be of any shape.

  • edge_feat (Tensor) - Edge featutes. The shape is \((N\_e,F\_e)\) where \(N\_e\) is the number of edges and \(F\_e\) is the number of edge features.

  • g (Graph) - The input graph.

Outputs:

  • Tensor, the output feature of shape \((N,D\_out)\) where \(N\) is the number of nodes and \(D\_out\) could be of any shape.

Raises

  • TypeError – if edge_embed type is not mindspore.nn.Cell or aggregator_type is not 'sum'.

  • TypeError – If in_feat_size or out_feat_size is not an int.

  • TypeError – If residual or bias is not a bool.

Supported Platforms:

Ascend GPU

Examples

>>> import mindspore as ms
>>> from mindspore_gl.nn import NNConv
>>> from mindspore_gl import GraphField
>>> n_nodes = 4
>>> n_edges = 7
>>> node_feat_size = 7
>>> edge_feat_size = 4
>>> src_idx = ms.Tensor([0, 1, 1, 2, 2, 3, 3], ms.int32)
>>> dst_idx = ms.Tensor([0, 0, 2, 1, 3, 0, 1], ms.int32)
>>> ones = ms.ops.Ones()
>>> node_feat = ones((n_nodes, node_feat_size), ms.float32)
>>> edge_feat = ones((n_edges, edge_feat_size), ms.float32)
>>> graph_field = GraphField(src_idx, dst_idx, n_nodes, n_edges)
>>> edge_func = ms.nn.Dense(edge_feat_size, 2)
>>> nnconv = NNConv(in_feat_size=node_feat_size, out_feat_size=2, edge_embed=edge_func)
>>> res = nnconv(node_feat, edge_feat, *graph_field.get_graph())
>>> print(res.shape)
(4, 2)