Change the output names list of YOLOv6 export to be dynamic (#112)

luxonis · Sep 24, 2024 · b6e79e1 · b6e79e1
1 parent 85db845
commit b6e79e1
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/yolo/export_yolov6.py b/yolo/export_yolov6.py
@@ -62,7 +62,7 @@ def export_onnx(self):
                         training=torch.onnx.TrainingMode.EVAL,
                         do_constant_folding=True,
                         input_names=['images'],
-                        output_names=['output1_yolov6r2', 'output2_yolov6r2', 'output3_yolov6r2'],
+                        output_names=[f"output{i+1}_yolov6r2" for i in range(self.num_branches)],
                         dynamic_axes=None)
 
         # check if the arhcitecture is correct

diff --git a/yolov6r3/yolo/export_gold_yolo.py b/yolov6r3/yolo/export_gold_yolo.py
@@ -52,7 +52,7 @@ def export_onnx(self):
                         training=torch.onnx.TrainingMode.EVAL,
                         do_constant_folding=True,
                         input_names=['images'],
-                        output_names=['output1_yolov6r2', 'output2_yolov6r2', 'output3_yolov6r2'],
+                        output_names=[f"output{i+1}_yolov6r2" for i in range(self.num_branches)],
                         dynamic_axes=None)
 
         # check if the arhcitecture is correct

diff --git a/yolov6r3/yolo/export_yolov6_r3.py b/yolov6r3/yolo/export_yolov6_r3.py
@@ -67,7 +67,7 @@ def export_onnx(self):
                         training=torch.onnx.TrainingMode.EVAL,
                         do_constant_folding=True,
                         input_names=['images'],
-                        output_names=['output1_yolov6r2', 'output2_yolov6r2', 'output3_yolov6r2'],
+                        output_names=[f"output{i+1}_yolov6r2" for i in range(self.num_branches)],
                         dynamic_axes=None)
 
         # check if the arhcitecture is correct