Add export output names specification for individual nodes (#112)

Co-authored-by: KlemenSkrlj <[email protected]>

configs/README.md

@@ -388,7 +388,7 @@ Here you can define configuration for exporting.
 | `mean_values`            | `list[float] \| None`             | `None`        | What mean values to use for input normalization. If not provided, inferred from augmentations  |
 | `upload_to_run`          | `bool`                            | `True`        | Whether to upload the exported files to tracked run as artifact                                |
 | `upload_url`             | `str \| None`                     | `None`        | Exported model will be uploaded to this URL if specified                                       |
-| `output_names`           | `list[str] \| None`               | `None`        | Optional list of output names to override the default ones                                     |
+| `output_names`           | `list[str] \| None`               | `None`        | Optional list of output names to override the default ones (deprecated)                        |
 
 ### `ONNX`
 
@@ -411,7 +411,6 @@ Option specific for `ONNX` export.
 
 ```yaml
 exporter:
-  output_names: ["output1", "output2"]
   onnx:
     opset_version: 11
   blobconverter:

configs/example_export.yaml

@@ -7,6 +7,8 @@ model:
     name: DetectionModel
     params:
       variant: light
+      head_params:
+       export_output_names: [output1_yolov6r2, output2_yolov6r2, output3_yolov6r2]
 
 loader:
   params:
@@ -42,7 +44,6 @@ trainer:
       last_epoch: -1
 
 exporter:
-  output_names: [output1_yolov6r2, output2_yolov6r2, output3_yolov6r2]
   onnx:
     opset_version: 11
   blobconverter:

luxonis_train/core/utils/archive_utils.py

@@ -186,26 +186,7 @@ def _get_head_outputs(
         if name == head_name:
             output_names.append(output["name"])
 
-    if output_names:
-        return output_names
-
-    # TODO: Fix this, will require refactoring custom ONNX output names
-    logger.error(
-        "ONNX model uses custom output names, trying to determine outputs based on the head type. "
-        "This will likely result in incorrect archive for multi-head models. "
-        "You can ignore this error if your model has only one head."
-    )
-
-    if head_type == "ClassificationHead":
-        return [outputs[0]["name"]]
-    elif head_type == "EfficientBBoxHead":
-        return [output["name"] for output in outputs]
-    elif head_type in ["SegmentationHead", "BiSeNetHead"]:
-        return [outputs[0]["name"]]
-    elif head_type == "EfficientKeypointBBoxHead":
-        return [outputs[0]["name"]]
-    else:
-        raise ValueError("Unknown head name")
+    return output_names
 
 
 def get_heads(
@@ -238,9 +219,15 @@ def get_heads(
                     task = str(next(iter(task.values())))
 
                 classes = _get_classes(node_name, task, class_dict)
-                head_outputs = _get_head_outputs(
-                    outputs, node_alias, node_name
-                )
+
+                export_output_names = nodes[node_alias].export_output_names
+                if export_output_names is not None:
+                    head_outputs = export_output_names
+                else:
+                    head_outputs = _get_head_outputs(
+                        outputs, node_alias, node_name
+                    )
+
                 if node_alias in head_names:
                     curr_head_name = f"{node_alias}_{len(head_names)}"  # add suffix if name is already present
                 else:

luxonis_train/models/luxonis_lightning.py

@@ -520,7 +520,38 @@ def export_onnx(self, save_path: str, **kwargs) -> list[str]:
             ]
         )
 
+        output_counts = defaultdict(int)
+        for node_name, outs in outputs.items():
+            output_counts[node_name] = sum(len(out) for out in outs.values())
+
         if self.cfg.exporter.output_names is not None:
+            logger.warning(
+                "The use of 'exporter.output_names' is deprecated and will be removed in a future version. "
+                "If 'node.export_output_names' are provided, they will take precedence and overwrite 'exporter.output_names'. "
+                "Please update your config to use 'node.export_output_names' directly."
+            )
+
+        export_output_names_used = False
+        export_output_names_dict = {}
+        for node_name, node in self.nodes.items():
+            if node.export_output_names is not None:
+                export_output_names_used = True
+                if len(node.export_output_names) != output_counts[node_name]:
+                    logger.warning(
+                        f"Number of provided output names for node {node_name} "
+                        f"({len(node.export_output_names)}) does not match "
+                        f"number of outputs ({output_counts[node_name]}). "
+                        f"Using default names."
+                    )
+                else:
+                    export_output_names_dict[node_name] = (
+                        node.export_output_names
+                    )
+
+        if (
+            not export_output_names_used
+            and self.cfg.exporter.output_names is not None
+        ):
             len_names = len(self.cfg.exporter.output_names)
             if len_names != len(output_order):
                 logger.warning(
@@ -529,18 +560,25 @@ def export_onnx(self, save_path: str, **kwargs) -> list[str]:
                 )
                 self.cfg.exporter.output_names = None
 
-        output_names = self.cfg.exporter.output_names or [
-            f"{node_name}/{output_name}/{i}"
-            for node_name, output_name, i in output_order
-        ]
+            output_names = self.cfg.exporter.output_names or [
+                f"{node_name}/{output_name}/{i}"
+                for node_name, output_name, i in output_order
+            ]
 
-        if not self.cfg.exporter.output_names:
-            idx = 1
-            # Set to output names required by DAI
-            for i, output_name in enumerate(output_names):
-                if output_name.startswith("EfficientBBoxHead"):
-                    output_names[i] = f"output{idx}_yolov6r2"
-                    idx += 1
+            if not self.cfg.exporter.output_names:
+                idx = 1
+                # Set to output names required by DAI
+                for i, output_name in enumerate(output_names):
+                    if output_name.startswith("EfficientBBoxHead"):
+                        output_names[i] = f"output{idx}_yolov6r2"
+                        idx += 1
+        else:
+            output_names = []
+            for node_name, output_name, i in output_order:
+                if node_name in export_output_names_dict:
+                    output_names.append(export_output_names_dict[node_name][i])
+                else:
+                    output_names.append(f"{node_name}/{output_name}/{i}")
 
         old_forward = self.forward
 

luxonis_train/nodes/base_node.py

@@ -119,6 +119,7 @@ def __init__(
         n_keypoints: int | None = None,
         in_sizes: Size | list[Size] | None = None,
         remove_on_export: bool = False,
+        export_output_names: list[str] | None = None,
         attach_index: AttachIndexType | None = None,
         _tasks: dict[TaskType, str] | None = None,
     ):
@@ -139,6 +140,11 @@ def __init__(
         @type in_sizes: Size | list[Size] | None
         @param in_sizes: List of input sizes for the node. Provide only
             in case the C{input_shapes} were not provided.
+        @type remove_on_export: bool
+        @param remove_on_export: If set to True, the node will be removed
+            from the model during export. Defaults to False.
+        @type export_output_names: list[str] | None
+        @param export_output_names: List of output names for the export.
         @type attach_index: AttachIndexType
         @param attach_index: Index of previous output that this node
             attaches to. Can be a single integer to specify a single
@@ -186,6 +192,7 @@ def __init__(
         self._n_keypoints = n_keypoints
         self._export = False
         self._remove_on_export = remove_on_export
+        self._export_output_names = export_output_names
         self._epoch = 0
         self._in_sizes = in_sizes
 
@@ -513,6 +520,11 @@ def remove_on_export(self) -> bool:
         """Getter for the remove_on_export attribute."""
         return self._remove_on_export
 
+    @property
+    def export_output_names(self) -> list[str] | None:
+        """Getter for the export_output_names attribute."""
+        return self._export_output_names
+
     def unwrap(self, inputs: list[Packet[Tensor]]) -> ForwardInputT:
         """Prepares inputs for the forward pass.
 

luxonis_train/nodes/heads/efficient_bbox_head.py

@@ -74,6 +74,22 @@ def __init__(
                 in_channels=self.in_channels[i],
             )
             self.heads.append(curr_head)
+        if (
+            self.export_output_names is None
+            or len(self.export_output_names) != self.n_heads
+        ):
+            if (
+                self.export_output_names is not None
+                and len(self.export_output_names) != self.n_heads
+            ):
+                logger.warning(
+                    f"Number of provided output names ({len(self.export_output_names)}) "
+                    f"does not match number of heads ({self.n_heads}). "
+                    f"Using default names."
+                )
+            self._export_output_names = [
+                f"output{i+1}_yolov6r2" for i in range(self.n_heads)
+            ]
 
     def forward(
         self, inputs: list[Tensor]

luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py

@@ -64,6 +64,8 @@ def __init__(
             for x in self.in_channels
         )
 
+        self._export_output_names = None
+
     def forward(
         self, inputs: list[Tensor]
     ) -> tuple[list[Tensor], list[Tensor], list[Tensor], list[Tensor]]:

tests/configs/archive_config.yaml

@@ -7,33 +7,30 @@ model:
     - name: EfficientBBoxHead
       inputs:
         - EfficientRep
+      params:
+        export_output_names: [bbox0, bbox1, bbox2]
 
     - name: EfficientKeypointBBoxHead
       inputs:
         - EfficientRep
+      params:
+        export_output_names: [effkpt0, effkpt1, effkpt2]
 
     - name: SegmentationHead
       inputs:
         - EfficientRep
+      params:
+        export_output_names: [seg0, seg1]
 
     - name: BiSeNetHead
       inputs:
         - EfficientRep
+      params:
+        export_output_names: [impl]
 
     - name: ClassificationHead
       inputs:
         - EfficientRep
-
-exporter:
-  output_names:
-    - seg0
-    - class0
-    - bbox0
-    - bbox1
-    - bbox2
-    - effkpt0
-    - effkpt1
-    - effkpt2
-    - impl
-    - seg1
+      params:
+        export_output_names: [class0]
 

tests/integration/parking_lot.json

@@ -47,7 +47,7 @@
                 "layout": "NCHW"
             },
             {
-                "name": "bbox-head/boundingbox/0",
+                "name": "output1_yolov6r2",
                 "dtype": "float32",
                 "shape": [
                     1,
@@ -58,7 +58,7 @@
                 "layout": "NCHW"
             },
             {
-                "name": "bbox-head/boundingbox/1",
+                "name": "output2_yolov6r2",
                 "dtype": "float32",
                 "shape": [
                     1,
@@ -69,7 +69,7 @@
                 "layout": "NCHW"
             },
             {
-                "name": "bbox-head/boundingbox/2",
+                "name": "output3_yolov6r2",
                 "dtype": "float32",
                 "shape": [
                     1,
@@ -164,9 +164,9 @@
                     "subtype": "yolov6"
                 },
                 "outputs": [
-                    "bbox-head/boundingbox/0",
-                    "bbox-head/boundingbox/1",
-                    "bbox-head/boundingbox/2"
+                    "output1_yolov6r2",
+                    "output2_yolov6r2",
+                    "output3_yolov6r2"
                 ]
             },
             {