mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge branch 'master' into exp

pull/12914/head
Ultralytics Assistant 2024-08-11 10:32:33 +02:00 committed by GitHub
parent a063108608 19ce90290c
commit 7da4888794
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
export.py
View File

@ -368,12 +368,11 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX
if simplify:
try:
cuda = torch.cuda.is_available()
check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1"))
import onnxsim
check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnxslim"))
import onnxslim
LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...")
model_onnx, check = onnxsim.simplify(model_onnx)
assert check, "assert check failed"
LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
model_onnx = onnxslim.slim(model_onnx)
onnx.save(model_onnx, f)
except Exception as e:
LOGGER.info(f"{prefix} simplifier failure: {e}")