mindspore_gl.nn.STConv

class mindspore_gl.nn.STConv(num_nodes: int, in_channels: int, hidden_channels: int, out_channels: int, kernel_size: int = 3, k: int = 3, bias: bool = True)[source]

Spatial-Temporal Graph Convolutional layer. From the paper A deep learning framework for traffic forecasting arXiv preprint arXiv:1709.04875, 2017. . The STGCN layer contains 2 temporal convolution layer and 1 graph convolution layer (ChebyNet).

Parameters

  • num_nodes (int) – number of nodes.

  • in_channels (int) – Input node feature size.

  • hidden_channels (int) – hidden feature size.

  • out_channels (int) – Output node feature size.

  • kernel_size (int, optional) – Convolutional kernel size. Default: 3.

  • k (int, optional) – Chebyshev filter size. Default: 3.

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

Inputs:

  • x (Tensor) - The input node features. The shape is \((B, T, N, (D_{in}))\) where \(B\) is the size of batch, \(T\) is the number of input time steps, \(N\) is the number of nodes, \((D_{in})\) should be equal to in_channels in Args.

  • edge_weight (Tensor) - Edge weights. The shape is \((N\_e,)\) where \(N\_e\) is the number of edges.

  • g (Graph) - The input graph.

Outputs:

  • Tensor, output node features with shape of \((B, D_{out}, N, T)\), where \(B\) is the size of batch, \((D_{out})\) should be the same as out_channels in Args, \(N\) is the number of nodes, \(T\) is the number of input time steps.

Raises

  • TypeError – If num_nodes or in_channels or out_channels or hidden_channels or kernel_size or is k not an int.

  • TypeError – If bias is not a bool.

Supported Platforms:

Ascend GPU

Examples

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore_gl.nn.temporal import STConv
>>> from mindspore_gl import GraphField
>>> from mindspore_gl.graph import norm
>>> n_nodes = 4
>>> n_edges = 6
>>> feat_size = 2
>>> edge_attr = ms.Tensor([1, 1, 1, 1, 1, 1], ms.float32)
>>> edge_index = ms.Tensor([[1, 1, 2, 2, 3, 3],
>>>                         [0, 2, 1, 3, 0, 1]], ms.int32)
>>> edge_index, edge_weight = norm(edge_index, n_nodes, edge_attr, 'sym')
>>> edge_weight = ms.ops.Reshape()(edge_weight, ms.ops.Shape()(edge_weight) + (1,))
>>> batch_size = 2
>>> input_time_steps = 5
>>> feat = ms.Tensor(np.ones((batch_size, input_time_steps, n_nodes, feat_size)), ms.float32)
>>> graph_field = GraphField(edge_index[0], edge_index[1], n_nodes, n_edges)
>>> stconv = STConv(num_nodes=n_nodes, in_channels=feat_size,
>>>                 hidden_channels=3, out_channels=2,
>>>                 kernel_size=2, k=2)
>>> out = stconv(feat, edge_weight, *graph_field.get_graph())
>>> print(out.shape)
(2, 3, 4, 2)