From bc9df47c11f0131b29f49ee2711d06219df557ae Mon Sep 17 00:00:00 2001 From: Sam_S Date: Sat, 5 Jun 2021 00:01:58 +0400 Subject: [PATCH] Add `output_names` argument for ONNX export with dynamic axes (#3456) * Add output names & dynamic axes for onnx export Add output_names and dynamic_axes names for all outputs in torch.onnx.export. The first four outputs of the model will have names output0, output1, output2, output3 * use first output only + cleanup Co-authored-by: Samridha Shrestha Co-authored-by: Glenn Jocher --- models/export.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/models/export.py b/models/export.py index 0d1147938e37..2db6a7699953 100644 --- a/models/export.py +++ b/models/export.py @@ -96,11 +96,14 @@ 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=opt.opset_version, input_names=['images'], + torch.onnx.export(model, img, f, verbose=False, opset_version=opt.opset_version, training=torch.onnx.TrainingMode.TRAINING if opt.train else torch.onnx.TrainingMode.EVAL, do_constant_folding=not opt.train, - 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) + input_names=['images'], + output_names=['output'], + dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640) + 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + } if opt.dynamic else None) # Checks model_onnx = onnx.load(f) # load onnx model