Update `dynamic` keys and comments for ONNX exports (#3772)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/3771/head
Yonghye Kwon 1 year ago committed by GitHub
parent 96cc28b439
commit e71c0e3748
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 4
      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

Loading…
Cancel
Save