graphnet.deployment.deployment_module module

Class(es) for deploying GraphNeT models in icetray as I3Modules.

class graphnet.deployment.deployment_module.DeploymentModule(model_config, state_dict, device, prediction_columns)[source]

Bases: Logger

Base DeploymentModule for GraphNeT.

Contains standard methods for loading models doing inference with them. Experiment-specific implementations may overwrite methods and should define __call__.

Construct DeploymentModule.

Parameters:

  • model_config (Union[ModelConfig, str]) – A model configuration file.

  • state_dict (Union[Dict[str, Tensor], str]) – A state dict for the model.

  • device (str, default: 'cpu') – The computational device to use. Defaults to “cpu”.

  • prediction_columns (Optional[List[str]], default: None) – Column names for each column in model output.