graphnet.deployment.icecube.inference_module module

IceCube I3InferenceModule.

Contains functionality for writing model predictions to i3 files.

class graphnet.deployment.icecube.inference_module.I3InferenceModule(pulsemap_extractor, model_config, state_dict, model_name, gcd_file, features, prediction_columns, pulsemap)[source]

Bases: DeploymentModule

General class for inference on i3 frames.

General class for inference on I3Frames (physics).

Parameters:

  • pulsemap_extractor (Union[List[I3FeatureExtractor], I3FeatureExtractor]) – The extractor used to extract the pulsemap.

  • model_config (Union[ModelConfig, str]) – The ModelConfig (or path to it) that summarizes the model used for inference.

  • state_dict (str) – Path to state_dict containing the learned weights.

  • model_name (str) – The name used for the model. Will help define the named entry in the I3Frame. E.g. “dynedge”.

  • gcd_file (str) – path to associated gcd file.

  • features (Optional[List[str]], default: None) – the features of the pulsemap that the model is expecting.

  • prediction_columns (Optional[List[str]], default: None) –

    column names for the predictions of the model. Will help define the named entry in the I3Frame.

    E.g. [‘energy_reco’]. Optional.

  • pulsemap (Optional[str], default: None) – the pulsmap that the model is expecting as input.