From b3029bd702bf17a001cfa504d2c3e668c9d36a87 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 28 Jun 2022 15:07:34 +0200 Subject: [PATCH 1/2] 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 --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index 72e170a30bf2..bfe55c3aad7e 100644 --- a/export.py +++ b/export.py @@ -119,8 +119,8 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst f = file.with_suffix('.onnx') torch.onnx.export( - model, - im, + model.cpu() if dynamic else model, # --dynamic only compatible with cpu + im.cpu() if dynamic else im, f, verbose=False, opset_version=opset, From 916f12ad4e18b8b8e5398cfedcc11e61dbce45b5 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 28 Jun 2022 15:14:56 +0200 Subject: [PATCH 2/2] Update export.py --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index bfe55c3aad7e..9daf39f871c2 100644 --- a/export.py +++ b/export.py @@ -499,8 +499,6 @@ def run( im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection # 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 for k, m in model.named_modules(): if isinstance(m, Detect): @@ -510,6 +508,8 @@ def run( for _ in range(2): 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 LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")