The Model Class

A module which wraps the given model into a Cython compatible object. Following the mapping procedure, the model is given additional attributes and functionality, which includes some minor adjustments to PyTorch Geometric.

AnalysisG.Model.Model(model = None):

model – A trainable model.

match_data_model_vars(Data sample) None

Inspects the input sample and creates an internal map where the inputs and outputs of the model are mapped to the variables of the sample.


Triggers the back-propagation of the learnable parameters.


Saves the current state of the model’s trainable parameters.


Restores a saved training state.

__call__(data) dict

Executes the model on the input data and outputs a dictionary with the following keys:

  • L_*: The loss associated with the particular output of the model.

  • A_*: The accuracy associated with the model’s prediction (if it is of classification type).

  • total: The total loss of the all the model outputs \mathcal{L}_{total} = \Sigma_{i} \mathcal{L} _i

  • graphs: Unpacks a batched data element with the output of the model’s output. The unpacking is a hacked method of the to_data_list() method of PyGeometric.

  • __param__ (dict) – A free parameter dictionary used to initialize the model with. This parameter relevant if the model has input parameters during the __init__ call.

  • train (bool) – Sets the model to training mode.

  • in_map (dict) – A mapping of the model’s input parameters and the training sample.

  • out_map (dict) – A mapping of the model’s output parameters to the truth training sample attributes.

  • loss_map (dict) – Returns the loss of each model output.

  • class_map (dict) – A map indicating which of the model’s output are used for classification modes.

  • code (Code) – Returns a Code object to rebuild the model object.

  • Epoch (int) – The epoch of the trained model.

  • KFold (int) – The kFold of the trained model.

  • Path (str) – The path of the model’s saved state.

  • RunName (str) – The run-name of the model.

  • model – Returns the original model type.

  • device (str) – The device that the model should be transferred to.

  • failure (bool) – Indicates whether the model’s code was not properly traced or initialized.

  • error_message (str) – Returns the error associated with the model.

  • KinematicMap (dict) – The kinematics to use for reconstructing the predicted mass from the nodes and topology.

Model Declarations (Example)

The Model wrapper class was introduced to do a preliminary scan of the model’s inputs and prevent the model from crashing due to missing features during training. To further improve performance during inference and training of a given model, the internal mapping only provides the model with the inputs it requires rather than injecting all features. To declare a specific input parameter, the pre-fixes G, N and E are used to indicate whether a Graph, Node or Edge feature is requested. A similar syntax is used to request truth features as inputs, by appending to the pre-fixes, _T.

class YourModel(...):

    def __init__(self):

        self.O_output = None
        self.L_output = "CEL"

        #... < some neural network initialization >...#

    def forward(self, edge_index, batch, G_<some graph feature>, N_<some node features>, E_<some edge feature>):

        # some logic...