Skip to content

Commit

Permalink
Update export.py for 2 dry runs (#2910)
Browse files Browse the repository at this point in the history
* Update export.py for 2 dry runs

* Update export.py
  • Loading branch information
glenn-jocher committed Apr 23, 2021
1 parent a2a514d commit 646386f
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions models/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 _ in range(2):
y = model(img) # dry runs

# TorchScript export -----------------------------------------------------------------------------------------------
prefix = colorstr('TorchScript:')
Expand All @@ -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)

Expand Down

0 comments on commit 646386f

Please sign in to comment.