Train-Only Heads

[JSabadin]

Luxonis Train - Remove on Export Feature

The remove_on_export option allows users to exclude specific nodes during model export. This is useful for auxiliary nodes that are not required in the final exported model.

Example Usage:

model:
  name: ddrnet_segmentation
  nodes:
    - name: DDRNet
      params:
        use_aux_heads: true
    - name: DDRNetSegmentationHead
      inputs: ["DDRNet"]
      alias: "segmentation_head"
      params:
           attach_index: -1
    - name: DDRNetSegmentationHead
      inputs: ["DDRNet"]
      alias: "aux_segmentation_head"
      params:
           attach_index: -2
           remove_on_export: true  # This node will be ignored during export

Nodes marked with remove_on_export: true will be skipped when exporting the model.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (dev@e16cde6). Learn more about missing BASE report.

✅ All tests successful. No failed tests found.

Additional details and impacted files

@@          Coverage Diff           @@
##             dev      #83   +/-   ##
======================================
  Coverage       ?   97.07%           
======================================
  Files          ?      140           
  Lines          ?     6120           
  Branches       ?        0           
======================================
  Hits           ?     5941           
  Misses         ?      179           
  Partials       ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kozlov721]

Left a small comment, otherwise LGTM.
However, it would be good to add a simple test case for a dummy network with a train-only head:

model = ...
results = model.eval()
# test results contain all heads
output_names = model.export_onnx()
# test output_names don't contain the train-only heads

[kozlov721]

I don't think the setter is necessary. There shouldn't be a need to change the value once it's set, and if a niche use case arises, we can rely on the underlying _remove_on_export.

IMHO, implementing a setter suggests that it's reasonable for the user to change the once-set value, which isn't the case here. Instead, making it accessible only via the original _remove_on_export attribute indicates that the value shouldn't be modified in regular use cases.

If we'd want to keep the setter anyway, I think implementing it as an actual setter:

@remove_on_export.setter
def remove_on_export(self, mode: bool) -> None:
    ...

is a bit nicer.

[klemen1999]

I agree, I don't think we need setter for this. I don't see it being changed after init in 99% of the cases

[klemen1999]

Thanks, LGTM!