From e71c0e3748c942b719e3dc9daa19cc2aaa5f4391 Mon Sep 17 00:00:00 2001 From: Yonghye Kwon Date: Mon, 17 Jul 2023 21:14:37 +0900 Subject: [PATCH] Update `dynamic` keys and comments for ONNX exports (#3772) Co-authored-by: Glenn Jocher --- ultralytics/engine/exporter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 7c7c779..828c98d 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -314,10 +314,10 @@ class Exporter: if dynamic: dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) if isinstance(self.model, SegmentationModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400) dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) elif isinstance(self.model, DetectionModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400) torch.onnx.export( self.model.cpu() if dynamic else self.model, # --dynamic only compatible with cpu