mlreco.models.layers.gnn.losses.node_primary module¶
-
class
mlreco.models.layers.gnn.losses.node_primary.NodePrimaryLoss(loss_config, batch_col=0, coords_col=(1, 4))[source]¶ Bases:
torch.nn.modules.module.ModuleTakes the two-channel node output of the GNN and optimizes node-wise scores such that nodes that initiate a particle cascade are given a high score (typically for showers).
For use in config: model:
name: cluster_gnn modules:
- grappa_loss:
- node_loss:
name: : primary batch_col : <column in the label data that specifies the batch ids of each voxel (default 3)> loss : <loss function: ‘CE’ or ‘MM’ (default ‘CE’)> reduction : <loss reduction method: ‘mean’ or ‘sum’ (default ‘sum’)> balance_classes : <balance loss per class: True or False (default False)> high_purity : <only penalize loss on groups with a single primary (default False)> use_group_pred : <redifines group ids according to edge predictions (default False)> group_pred_alg : <algorithm used to predict cluster labels: ‘threshold’ or ‘score’ (default ‘score’)>
-
__init__(loss_config, batch_col=0, coords_col=(1, 4))[source]¶ Initializes internal Module state, shared by both nn.Module and ScriptModule.
-
__module__= 'mlreco.models.layers.gnn.losses.node_primary'¶
-
forward(out, clusters)[source]¶ Applies the requested loss on the node prediction.
- Parameters
out (dict) – ‘node_pred’ (torch.tensor) : (C,2) Two-channel node predictions ‘clusts’ ([np.ndarray]) : [(N_0), (N_1), …, (N_C)] Cluster ids (‘edge_pred’ (torch.tensor): (C,2) Two-channel edge predictions, optional) (‘edge_index’ (np.ndarray) : (E,2) Incidence matrix, optional)
clusters ([torch.tensor]) – (N,8) [x, y, z, batchid, value, id, groupid, shape]
- Returns
loss, accuracy, cluster count
- Return type
double
-
training: bool¶