From 19ce90290c6bd013b69dbe037e8ef1bc3d3c32d0 Mon Sep 17 00:00:00 2001 From: Kayzwer <68285002+Kayzwer@users.noreply.github.com> Date: Sun, 11 Aug 2024 16:31:53 +0800 Subject: [PATCH] Replace ONNXSim with ONNXSlim (#13236) Signed-off-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com> --- export.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/export.py b/export.py index c1524ec2d48a..dfb1c06fb5e2 100644 --- a/export.py +++ b/export.py @@ -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}")