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

Merge branch 'master' into master

pull/12599/head
Ultralytics Assistant 2024-08-11 10:32:24 +02:00 committed by GitHub
parent f811b6a58a 19ce90290c
commit 83cd15345a
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: if simplify:
try: try:
cuda = torch.cuda.is_available() cuda = torch.cuda.is_available()
check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1")) check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnxslim"))
import onnxsim import onnxslim
LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...") LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
model_onnx, check = onnxsim.simplify(model_onnx) model_onnx = onnxslim.slim(model_onnx)
assert check, "assert check failed"
onnx.save(model_onnx, f) onnx.save(model_onnx, f)
except Exception as e: except Exception as e:
LOGGER.info(f"{prefix} simplifier failure: {e}") LOGGER.info(f"{prefix} simplifier failure: {e}")