From b6e79e1261d7e4743ee36cff969d50ab0f635b30 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20=C4=8Cuhel?=
 <79118988+HonzaCuhel@users.noreply.github.com>
Date: Tue, 24 Sep 2024 11:06:45 +0200
Subject: [PATCH] Change the output names list of YOLOv6 export to be dynamic
 (#112)

---
 yolo/export_yolov6.py             | 2 +-
 yolov6r3/yolo/export_gold_yolo.py | 2 +-
 yolov6r3/yolo/export_yolov6_r3.py | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/yolo/export_yolov6.py b/yolo/export_yolov6.py
index 5beeb42..2f0a286 100644
--- 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
index cd82c47..68b1e2e 100644
--- 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
index 30192dc..5a22444 100644
--- 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