dymad.modules.gnn

Classes

GNN(input_dim, hidden_dim, output_dim, ...)

Configurable Graph Neural Network using a choice of GCL (e.g., SAGEConv, ChebConv) and activations.

IdenCatGNN(input_dim, hidden_dim, ...[, ...])

Identity concatenation GNN.

ResBlockGNN(input_dim, hidden_dim, ...[, ...])

Residual block with GNN as the nonlinearity.
class dymad.modules.gnn.GNN(input_dim, hidden_dim, output_dim, n_layers, *, gcl='sage', gcl_opts=None, activation='prelu', weight_init='xavier_uniform', bias_init='zeros', gain=1.0, end_activation=True, dtype=None, device=None)

Bases: Module

Configurable Graph Neural Network using a choice of GCL (e.g., SAGEConv, ChebConv) and activations.

Due to the implementation, the GNN is applied sequentially to batch data.

To interface with other parts of the code, the model assumes the input to be node-wise, (…, n_nodes, n_input), but the output is reshaped to concatenate features across nodes, (…, n_nodes * n_output). See forward method for details.

Parameters:

  • input_dim (int) – Dimension of input node features.

  • hidden_dim (int) – Dimension of hidden layers.

  • output_dim (int) – Dimension of output node features.

  • n_layers (int) – Number of GCL layers.

  • gcl (str | nn.Module | type, default='sage') – Graph convolution layer type or instance.

  • gcl_opts (dict, default={}) – Options passed to the GCL constructor.

  • activation (str | nn.Module | type, default='prelu') – Activation function.

  • weight_init (str | callable, default='xavier_uniform') – Weight initializer.

  • bias_init (str | callable, default='zeros') – Bias initializer.

  • gain (float, default=1.0) – Extra gain modifier for weight initialization.

  • end_activation (bool, default=True) – Whether to apply activation after last layer.

diagnostic_info()
Return type:

str

forward(x, edge_index, edge_weights, edge_attr, **kwargs)

Forward pass through the GNN.

  • x (…, n_nodes, n_features).

  • edge_index (…, n_edges, 2).

  • edge_weights (…, n_edges).

  • edge_attr (…, n_edges, n_edge_features).

  • Returns (…, n_nodes*n_new_features).

If …=1, we can process the entire batch in one go. Otherwise, we aggregate the graph on the fly so the shapes are reduced to the first case. The aggregation takes a bit more time.

class dymad.modules.gnn.IdenCatGNN(input_dim, hidden_dim, output_dim, n_layers, gcl='sage', gcl_opts=None, activation='prelu', weight_init='xavier_uniform', bias_init='zeros', gain=1.0, end_activation=True, dtype=None, device=None)

Bases: GNN

Identity concatenation GNN.

This GNN concatenates the input with the output of the GNN.

Note

The output dimension represents the total output features and must be greater than the input dimension.

See GNN for the arguments.

forward(x, edge_index, edge_weights, edge_attr, **kwargs)

Forward pass through the GNN.

  • x (…, n_nodes, n_features).

  • edge_index (…, n_edges, 2).

  • edge_weights (…, n_edges).

  • edge_attr (…, n_edges, n_edge_features).

  • Returns (…, n_nodes*n_new_features).

If …=1, we can process the entire batch in one go. Otherwise, we aggregate the graph on the fly so the shapes are reduced to the first case. The aggregation takes a bit more time.

class dymad.modules.gnn.ResBlockGNN(input_dim, hidden_dim, output_dim, n_layers, gcl='sage', gcl_opts=None, activation='prelu', weight_init='xavier_uniform', bias_init='zeros', gain=1.0, end_activation=True, dtype=None, device=None)

Bases: GNN

Residual block with GNN as the nonlinearity.

See GNN for the arguments.

forward(x, edge_index, edge_weights, edge_attr, **kwargs)

Forward pass through the GNN.

  • x (…, n_nodes, n_features).

  • edge_index (…, n_edges, 2).

  • edge_weights (…, n_edges).

  • edge_attr (…, n_edges, n_edge_features).

  • Returns (…, n_nodes*n_new_features).

If …=1, we can process the entire batch in one go. Otherwise, we aggregate the graph on the fly so the shapes are reduced to the first case. The aggregation takes a bit more time.