mlreco.models.layers.gnn.message_passing.meta module

class mlreco.models.layers.gnn.message_passing.meta.MetaLayerModel(cfg)[source]

Bases: torch.nn.modules.module.Module

MetaLayer GNN Module for extracting node/edge/global features

__init__(cfg)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(node_features, edge_indices, edge_features, xbatch)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

__module__ = 'mlreco.models.layers.gnn.message_passing.meta'
training: bool
class mlreco.models.layers.gnn.message_passing.meta.EdgeLayer(node_in, edge_in, edge_out, leakiness=0.0)[source]

Bases: torch.nn.modules.module.Module

An EdgeModel for predicting edge features.

Example: Parent-Child Edge prediction and EM primary assignment prediction.

INPUTS:

DEFINITIONS:

E: number of edges F_x: number of node features F_e: number of edge features F_u: number of global features F_o: number of output edge features B: number of graphs (same as batch size)

If an entry i->j is an edge, then we have source node feature F^i_x, target node feature F^j_x, and edge features F_e.

  • source: [E, F_x] Tensor, where E is the number of edges

  • target: [E, F_x] Tensor, where E is the number of edges

  • edge_attr: [E, F_e] Tensor, indicating input edge features.

  • global_features: [B, F_u] Tensor, where B is the number of graphs

(equivalent to number of batches).

  • batch: [E] Tensor containing batch indices for each edge from 0 to B-1.

Returns

[E, F_o] Tensor with F_o output edge features.

Return type

  • output

__init__(node_in, edge_in, edge_out, leakiness=0.0)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(src, dest, edge_attr, u=None, batch=None)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

__module__ = 'mlreco.models.layers.gnn.message_passing.meta'
training: bool
class mlreco.models.layers.gnn.message_passing.meta.NodeLayer(node_in, node_out, edge_in, leakiness=0.0, reduction='mean')[source]

Bases: torch.nn.modules.module.Module

NodeModel for node feature prediction.

Example: Particle Classification using node-level features.

INPUTS:

DEFINITIONS:

N: number of nodes F_x: number of node features F_e: number of edge features F_u: number of global features F_o: number of output node features B: number of graphs (same as batch size)

If an entry i->j is an edge, then we have source node feature F^i_x, target node feature F^j_x, and edge features F_e.

  • source: [E, F_x] Tensor, where E is the number of edges

  • target: [E, F_x] Tensor, where E is the number of edges

  • edge_attr: [E, F_e] Tensor, indicating input edge features.

  • global_features: [B, F_u] Tensor, where B is the number of graphs

(equivalent to number of batches).

  • batch: [E] Tensor containing batch indices for each edge from 0 to B-1.

Returns

[C, F_o] Tensor with F_o output node feature

Return type

  • output

__init__(node_in, node_out, edge_in, leakiness=0.0, reduction='mean')[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x, edge_index, edge_attr, u, batch)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

__module__ = 'mlreco.models.layers.gnn.message_passing.meta'
training: bool
class mlreco.models.layers.gnn.message_passing.meta.GlobalModel(node_in, batch_size, global_out, leakiness=0.0, reduction='mean')[source]

Bases: torch.nn.modules.module.Module

Global Model for global feature prediction.

Example: event classification (graph classification) over the whole image within a batch.

Do Hierarchical Pooling to reduce features

INPUTS:

DEFINITIONS:

N: number of nodes F_x: number of node features F_e: number of edge features F_u: number of global features F_o: number of output node features B: number of graphs (same as batch size)

If an entry i->j is an edge, then we have source node feature F^i_x, target node feature F^j_x, and edge features F_e.

  • source: [E, F_x] Tensor, where E is the number of edges

  • target: [E, F_x] Tensor, where E is the number of edges

  • edge_attr: [E, F_e] Tensor, indicating input edge features.

  • global_features: [B, F_u] Tensor, where B is the number of graphs

(equivalent to number of batches).

  • batch: [E] Tensor containing batch indices for each edge from 0 to B-1.

Returns

[C, F_o] Tensor with F_o output node feature

Return type

  • output

__init__(node_in, batch_size, global_out, leakiness=0.0, reduction='mean')[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x, edge_index, edge_attr, u, batch)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

__module__ = 'mlreco.models.layers.gnn.message_passing.meta'
training: bool
mlreco.models.layers.gnn.message_passing.gatconv module mlreco.models.layers.gnn.message_passing.nnconv module