瀏覽代碼

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 年之前
父節點
當前提交
0c1324067c
沒有發現已知的金鑰在資料庫的簽署中 GPG 金鑰 ID: 4AEE18F83AFDEB23
共有 1 個檔案被更改,包括 4 行新增4 行删除
  1. +4
    -4
      export.py

+ 4
- 4
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,
@@ -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)")


Loading…
取消
儲存