From b32c72df29d187a24317908fb45a24828ca69790 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 24 Apr 2021 00:09:50 +0200 Subject: [PATCH 1/2] Update export.py for 2 dry runs --- models/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/export.py b/models/export.py index 5b7b6bda08ae..023e29fc136c 100644 --- a/models/export.py +++ b/models/export.py @@ -58,7 +58,8 @@ # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) model.model[-1].export = not opt.grid # set Detect() layer grid export - y = model(img) # dry run + for i in range(2): + y = model(img) # dry runs # TorchScript export ----------------------------------------------------------------------------------------------- prefix = colorstr('TorchScript:') @@ -80,7 +81,6 @@ print(f'{prefix} starting export with onnx {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'], 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) From c4d94136bc10f519ee4a24a2644960544820536d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 24 Apr 2021 00:10:17 +0200 Subject: [PATCH 2/2] Update export.py --- models/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/export.py b/models/export.py index 023e29fc136c..a820ff94ac17 100644 --- a/models/export.py +++ b/models/export.py @@ -58,7 +58,7 @@ # elif isinstance(m, models.yolo.Detect): # m.forward = m.forward_export # assign forward (optional) model.model[-1].export = not opt.grid # set Detect() layer grid export - for i in range(2): + for _ in range(2): y = model(img) # dry runs # TorchScript export -----------------------------------------------------------------------------------------------