Browse Source

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
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
0c1324067c
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
  1. +4
    -4
      export.py

+ 4
- 4
export.py View File

f = file.with_suffix('.onnx') f = file.with_suffix('.onnx')


torch.onnx.export( torch.onnx.export(
model,
im,
model.cpu() if dynamic else model, # --dynamic only compatible with cpu
im.cpu() if dynamic else im,
f, f,
verbose=False, verbose=False,
opset_version=opset, opset_version=opset,
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):


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)")



Loading…
Cancel
Save