From c265dcf697b2665da662150ab75a6c592c4fe130 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Fri, 12 Feb 2021 22:23:52 +0530 Subject: [PATCH 1/6] add: dynamic onnx export --- detect_onnx.py | 12 ++++++++++++ models/export.py | 11 +++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 detect_onnx.py diff --git a/detect_onnx.py b/detect_onnx.py new file mode 100644 index 000000000000..af2e1d65dede --- /dev/null +++ b/detect_onnx.py @@ -0,0 +1,12 @@ +import numpy as np +import onnxruntime + +x = np.zeros((1, 3, 320, 320)).astype('float32') + +session = onnxruntime.InferenceSession("yolov5s.onnx", None) +input_name = session.get_inputs()[0].name +output_name = session.get_outputs()[0].name +print(input_name, output_name) + +result = session.run([output_name], {input_name: x}) +print(result[0].shape) \ No newline at end of file diff --git a/models/export.py b/models/export.py index 057658af53dc..e4906be2d706 100644 --- a/models/export.py +++ b/models/export.py @@ -22,6 +22,7 @@ parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width + parser.add_argument('--dynamic', action='store_true', help='dynamic onnx export') parser.add_argument('--batch-size', type=int, default=1, help='batch size') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand @@ -69,8 +70,14 @@ print('\nStarting ONNX export with onnx %s...' % onnx.__version__) f = opt.weights.replace('.pt', '.onnx') # filename - torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], - output_names=['classes', 'boxes'] if y is None else ['output']) + if opt.dynamic: + torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], + output_names=['output'], dynamic_axes={'images': {0: 'batch_size', 2: 'image_width', + 3: 'image_height'}, + 'output': {0: 'batch_size', 2: 'op1', 3: 'op2'}}) + else: + torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], + output_names=['classes', 'boxes'] if y is None else ['output']) # Checks onnx_model = onnx.load(f) # load onnx model From b8b5a86d9569f651eff2906ddde97b8f379093b3 Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Sat, 13 Feb 2021 20:23:41 +0530 Subject: [PATCH 2/6] delete: test onnx inference --- detect_onnx.py | 12 ------------ 1 file changed, 12 deletions(-) delete mode 100644 detect_onnx.py diff --git a/detect_onnx.py b/detect_onnx.py deleted file mode 100644 index af2e1d65dede..000000000000 --- a/detect_onnx.py +++ /dev/null @@ -1,12 +0,0 @@ -import numpy as np -import onnxruntime - -x = np.zeros((1, 3, 320, 320)).astype('float32') - -session = onnxruntime.InferenceSession("yolov5s.onnx", None) -input_name = session.get_inputs()[0].name -output_name = session.get_outputs()[0].name -print(input_name, output_name) - -result = session.run([output_name], {input_name: x}) -print(result[0].shape) \ No newline at end of file From c453bb86c5bc8a4459529ef4490b96a93c3d3fe8 Mon Sep 17 00:00:00 2001 From: Shivam Swanrkar Date: Sun, 14 Feb 2021 14:16:36 +0530 Subject: [PATCH 3/6] fix dynamic output axis --- models/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index e4906be2d706..e8d8589dafa0 100644 --- a/models/export.py +++ b/models/export.py @@ -74,7 +74,7 @@ torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], output_names=['output'], dynamic_axes={'images': {0: 'batch_size', 2: 'image_width', 3: 'image_height'}, - 'output': {0: 'batch_size', 2: 'op1', 3: 'op2'}}) + 'output': {0: 'batch_size', 1: 'op1', 2: 'op2'}}) else: torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], output_names=['classes', 'boxes'] if y is None else ['output']) From 5ec68db3a71cebdb2db44eda8c3dcf80da6458ff Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 15 Feb 2021 20:58:31 -0800 Subject: [PATCH 4/6] Code reduction --- models/export.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/models/export.py b/models/export.py index e8d8589dafa0..96d871e2ddb0 100644 --- a/models/export.py +++ b/models/export.py @@ -70,14 +70,10 @@ print('\nStarting ONNX export with onnx %s...' % onnx.__version__) f = opt.weights.replace('.pt', '.onnx') # filename - if opt.dynamic: - torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], - output_names=['output'], dynamic_axes={'images': {0: 'batch_size', 2: 'image_width', - 3: 'image_height'}, - 'output': {0: 'batch_size', 1: 'op1', 2: 'op2'}}) - else: - torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], - output_names=['classes', 'boxes'] if y is None else ['output']) + torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], + output_names=['classes', 'boxes'] if y is None else ['output'], + dynamic_axes={'images': {0: 'batch_size', 2: 'image_width', 3: 'image_height'}, + 'output': {0: 'batch_size', 1: 'op1', 2: 'op2'}} if opt.dynamic else None) # Checks onnx_model = onnx.load(f) # load onnx model From e7e4a59d10ef3663cdd0d376aa3b41d725927b8e Mon Sep 17 00:00:00 2001 From: Aditya Lohia Date: Sun, 21 Feb 2021 13:15:10 +0530 Subject: [PATCH 5/6] fix: dynamic output axes, dynamic input naming --- models/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/export.py b/models/export.py index 96d871e2ddb0..c9ce17215753 100644 --- a/models/export.py +++ b/models/export.py @@ -72,8 +72,8 @@ f = opt.weights.replace('.pt', '.onnx') # filename torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], output_names=['classes', 'boxes'] if y is None else ['output'], - dynamic_axes={'images': {0: 'batch_size', 2: 'image_width', 3: 'image_height'}, - 'output': {0: 'batch_size', 1: 'op1', 2: 'op2'}} if opt.dynamic else None) + dynamic_axes={'images': {0: 'batch_size', 2: 'image_height', 3: 'image_width'}, + 'output': {0: 'batch_size', 1: 'anchors', 2: 'grid_y', 3: 'grid_x', 4: 'output'}} if opt.dynamic else None) # Checks onnx_model = onnx.load(f) # load onnx model From d6676f34f285027985a4e9b3f498ab7748fa0335 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 21 Feb 2021 21:46:30 -0800 Subject: [PATCH 6/6] Remove fixed axes --- models/export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/models/export.py b/models/export.py index c9ce17215753..cc817871f218 100644 --- a/models/export.py +++ b/models/export.py @@ -22,7 +22,7 @@ parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/ parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width - parser.add_argument('--dynamic', action='store_true', help='dynamic onnx export') + parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') parser.add_argument('--batch-size', type=int, default=1, help='batch size') opt = parser.parse_args() opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand @@ -72,8 +72,8 @@ f = opt.weights.replace('.pt', '.onnx') # filename torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], output_names=['classes', 'boxes'] if y is None else ['output'], - dynamic_axes={'images': {0: 'batch_size', 2: 'image_height', 3: 'image_width'}, - 'output': {0: 'batch_size', 1: 'anchors', 2: 'grid_y', 3: 'grid_x', 4: 'output'}} if opt.dynamic else None) + dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640) + 'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None) # Checks onnx_model = onnx.load(f) # load onnx model