glenn-jocher commited on
Commit
646386f
·
unverified ·
1 Parent(s): a2a514d

Update export.py for 2 dry runs (#2910)

Browse files

* Update export.py for 2 dry runs

* Update export.py

Files changed (1) hide show
  1. models/export.py +2 -2
models/export.py CHANGED
@@ -58,7 +58,8 @@ if __name__ == '__main__':
58
  # elif isinstance(m, models.yolo.Detect):
59
  # m.forward = m.forward_export # assign forward (optional)
60
  model.model[-1].export = not opt.grid # set Detect() layer grid export
61
- y = model(img) # dry run
 
62
 
63
  # TorchScript export -----------------------------------------------------------------------------------------------
64
  prefix = colorstr('TorchScript:')
@@ -80,7 +81,6 @@ if __name__ == '__main__':
80
  print(f'{prefix} starting export with onnx {onnx.__version__}...')
81
  f = opt.weights.replace('.pt', '.onnx') # filename
82
  torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
83
- output_names=['classes', 'boxes'] if y is None else ['output'],
84
  dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
85
  'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
86
 
 
58
  # elif isinstance(m, models.yolo.Detect):
59
  # m.forward = m.forward_export # assign forward (optional)
60
  model.model[-1].export = not opt.grid # set Detect() layer grid export
61
+ for _ in range(2):
62
+ y = model(img) # dry runs
63
 
64
  # TorchScript export -----------------------------------------------------------------------------------------------
65
  prefix = colorstr('TorchScript:')
 
81
  print(f'{prefix} starting export with onnx {onnx.__version__}...')
82
  f = opt.weights.replace('.pt', '.onnx') # filename
83
  torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
 
84
  dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
85
  'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
86