zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Fix ONNX `--dynamic` export on GPU (#8378) 

* Fix ONNX `--dynamic` export on GPU

Patch forces --dynamic export model and image to CPU. Resolves bug raised in https://github.com/ultralytics/yolov5/issues/8377

* Update export.py
This commit is contained in:
Glenn Jocher 2022-06-28 15:22:15 +02:00 committed by GitHub
parent 50ff6eee31
commit 0c1324067c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
export.py
View File

@ -119,8 +119,8 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
f = file.with_suffix('.onnx') f = file.with_suffix('.onnx')
torch.onnx.export( torch.onnx.export(
model, model.cpu() if dynamic else model, # --dynamic only compatible with cpu
im, im.cpu() if dynamic else im,
f, f,
verbose=False, verbose=False,
opset_version=opset, opset_version=opset,
@ -499,8 +499,6 @@ def run(
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
# Update model # Update model
if half and not coreml and not xml:
im, model = im.half(), model.half() # to FP16
model.train() if train else model.eval() # training mode = no Detect() layer grid construction model.train() if train else model.eval() # training mode = no Detect() layer grid construction
for k, m in model.named_modules(): for k, m in model.named_modules():
if isinstance(m, Detect): if isinstance(m, Detect):
@ -510,6 +508,8 @@ def run(
for _ in range(2): for _ in range(2):
y = model(im) # dry runs y = model(im) # dry runs
if half and not coreml:
im, model = im.half(), model.half() # to FP16
shape = tuple(y[0].shape) # model output shape shape = tuple(y[0].shape) # model output shape
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")