Merge remote-tracking branch 'maxwell_fork/both_models' into multitask_bundle_combine

Author: Aske-Rosted

src/graphnet/deployment/deployment_module.py

@@ -111,7 +111,8 @@ def _resolve_prediction_columns(
                             prediction_columns[i]
                         )
                 else:
-                    resolved_prediction_columns.append(model.prediction_labels)
+                    # Only Take First Label
+                    resolved_prediction_columns.append(model.prediction_labels[0])
             return resolved_prediction_columns
 
     def _inference(self, data: Union[Data, Batch]) -> List[np.ndarray]:
@@ -139,5 +140,9 @@ def _inference(self, data: Union[Data, Batch]) -> List[np.ndarray]:
                 output = model(data=data[_])
                 for k in range(len(output)):
                     output[k] = output[k].detach().cpu().numpy()
+                if output[0].shape[1] > 1:
+                    output = np.delete(output[0], 1, axis=1)
+                else:
+                    output = output[0]
                 outputs.append(output)
             return outputs

src/graphnet/deployment/icecube/__init__.py

@@ -1,5 +1,5 @@
 """Deployment modules specific to IceCube."""
 
-from .inference_module import I3InferenceModule, I3ParticleInferenceModule
+from .inference_module import I3InferenceModule, I3ParticleInferenceModule, I3MultipleModelInferenceModule
 from .cleaning_module import I3PulseCleanerModule
 from .i3deployer import I3Deployer

src/graphnet/deployment/icecube/inference_module.py

@@ -4,10 +4,8 @@
 """
 
 from typing import List, Union, Optional, TYPE_CHECKING, Dict, Any
-
 import numpy as np
 from torch_geometric.data import Data, Batch
-
 from torch.cuda import OutOfMemoryError
 from time import sleep
 import torch
@@ -51,7 +49,6 @@ def __init__(
         prediction_columns: Optional[Union[List[str], None]] = None,
         pulsemap: Optional[str] = None,
         multiple_models: bool = False,
-        key_name: Optional[str] = None,
         requirements: Optional[callable] = None,
         device: Optional[str] = "cpu",
         batch_size: Optional[int] = 1,
@@ -75,8 +72,6 @@ def __init__(
                                 E.g. ['energy_reco']. Optional.
             pulsemap: the pulsmap that the model is expecting as input.
             multiple_models: process multiple models with the same feature set at once.
-            key_name: The name used for the key in the I3Frame. Will help define the
-                     named entry in the I3Frame. E.g. "dynedge_predictions".
         """
         super().__init__(
             model_config=model_config,
@@ -130,7 +125,6 @@ def __init__(
         self._num_threads = num_threads
         self._inference_speed_check = inference_speed_check
         self._multiple_models = multiple_models
-        self._key_name = key_name
         # Set GCD file for pulsemap extractor
         if gcd_file is not None:
             for i3_extractor in self._i3_extractors:
@@ -266,26 +260,31 @@ def _create_dictionary(
     ) -> Dict[str, Any]:
         """Transform predictions into a dictionary."""
         data = {}
-        for i in range(dim):
-            data[self.model_name + "_" + self.prediction_columns[i]] = (
-                I3Double(float(predictions[i]))
-            )
+        if self._multiple_models == True:
+            for i, key in enumerate(self.model_name):
+                data[key] = float(predictions[i])
+        else:
+            for i in range(dim):
+                data[self.model_name + "_" + self.prediction_columns[i]] = (
+                    I3Double(float(predictions[i]))
+                )
 
-            # try:
-            #     assert len(predictions[:, i]) == 1
-            #     data[
-            #         self.model_name + "_" + self.prediction_columns[i]
-            #     ] = I3Double(float(predictions[:, i][0]))
-            # except IndexError:
-            #     data[
-            #         self.model_name + "_" + self.prediction_columns[i]
-            #     ] = I3Double(predictions[0])
+                # try:
+                #     assert len(predictions[:, i]) == 1
+                #     data[
+                #         self.model_name + "_" + self.prediction_columns[i]
+                #     ] = I3Double(float(predictions[:, i][0]))
+                # except IndexError:
+                #     data[
+                #         self.model_name + "_" + self.prediction_columns[i]
+                #     ] = I3Double(predictions[0])
         return data
 
     def _apply_model(self, data: Data) -> np.ndarray:
         """Apply model to `Data` and case-handling."""
         if data is not None:
             predictions = self._inference(data)
+            #print(predictions, type(predictions), type(predictions[0]))
             if isinstance(predictions, list):
                 predictions = np.concatenate(
                     [pred.flatten() for pred in predictions]
@@ -479,3 +478,25 @@ def _add_to_frame(self, frame, data):
         # for all the other values in data, add them to an I3Dictionary
         super()._add_to_frame(frame=frame, data=data)
         return
+
+class I3MultipleModelInferenceModule(I3InferenceModule):
+    """I3InferenceModule for I3Particle data."""
+
+    def __init__(
+        self,
+        key_name: List[str],
+        **kwargs,
+    ):
+        """
+        key_name: The name used for the key in the I3Frame. Will help define the
+                     named entry in the I3Frame.
+        """
+        super().__init__(**kwargs)
+
+        self._key_name = key_name
+
+    def _add_to_frame(self, frame, data):
+
+        i3_score_container = dataclasses.I3MapStringDouble(data)
+        frame.Put(self._key_name, i3_score_container)
+

src/graphnet/models/standard_model.py

@@ -94,18 +94,20 @@ def __init__(
         self._data_representation = data_representation
         self.backbone = backbone
         self._split = split
-        assert (
-            sum(self._split[0]) == self.backbone.nb_outputs
-        ), "Split dimensions do not match backbone output dimension check your configuration"
-
-        if learned_multitask_weights != -1:
-            assert isinstance(tasks, list)
-            # init the module for learned task weights
-            self.loss_weight_balancing = LossWeightBalancing(
-                tasks, late_activation=learned_multitask_weights
-            )
-        else:
-            self.loss_weight_balancing = None
+
+        if self._split is not None:
+            assert (
+                sum(self._split[0]) == self.backbone.nb_outputs
+            ), "Split dimensions do not match backbone output dimension check your configuration"
+
+            if learned_multitask_weights != -1:
+                assert isinstance(tasks, list)
+                # init the module for learned task weights
+                self.loss_weight_balancing = LossWeightBalancing(
+                    tasks, late_activation=learned_multitask_weights
+                )
+            else:
+                self.loss_weight_balancing = None
 
     def compute_loss(
         self, preds: Tensor, data: List[Data], verbose: bool = False