Update export.py for 2 dry runs (#2910)
Browse files* Update export.py for 2 dry runs
* Update export.py
- 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 |
-
|
|
|
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 |
|