mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Replace ONNXSim with ONNXSlim (#13236)

Browse Source 
Signed-off-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
This commit is contained in:
Kayzwer 2024-08-11 16:31:53 +08:00 committed by GitHub
parent 6096750fcc
commit 19ce90290c
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}")