HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.27 replace onnxsim with onnxslim (#12989)

Browse source 
Co-authored-by: inisis <desmond.yao@buaa.edu.cn>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: inisis <46103969+inisis@users.noreply.github.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Kayzwer 2024-06-01 01:03:34 +08:00 • committed by GitHub
parent dd13707bf8
commit 8fb140688a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 27 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

24
ultralytics/engine/exporter.py
View file

@ -384,7 +384,7 @@ class Exporter:
"""YOLOv8 ONNX export."""
requirements = ["onnx>=1.12.0"]
if self.args.simplify:
requirements += ["cmake", "onnxsim>=0.4.33", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
requirements += ["onnxslim==0.1.28", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
check_requirements(requirements)
import onnx # noqa
@ -421,14 +421,17 @@ class Exporter:
# Simplify
if self.args.simplify:
try:
import onnxsim
import onnxslim
LOGGER.info(f"{prefix} simplifying with onnxsim {onnxsim.__version__}...")
# subprocess.run(f'onnxsim "{f}" "{f}"', shell=True)
model_onnx, check = onnxsim.simplify(model_onnx)
assert check, "Simplified ONNX model could not be validated"
LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
model_onnx = onnxslim.slim(model_onnx)
# ONNX Simplifier (deprecated as must be compiled with 'cmake' in aarch64 and Conda CI environments)
# import onnxsim
# model_onnx, check = onnxsim.simplify(model_onnx)
# assert check, "Simplified ONNX model could not be validated"
except Exception as e:
LOGGER.info(f"{prefix} simplifier failure: {e}")
LOGGER.warning(f"{prefix} simplifier failure: {e}")
# Metadata
for k, v in self.metadata.items():
@ -672,8 +675,8 @@ class Exporter:
def export_engine(self, prefix=colorstr("TensorRT:")):
"""YOLOv8 TensorRT export https://developer.nvidia.com/tensorrt."""
assert self.im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. use 'device=0'"
self.args.simplify = True
f_onnx, _ = self.export_onnx() # run before trt import https://github.com/ultralytics/ultralytics/issues/7016
# self.args.simplify = True
f_onnx, _ = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
try:
import tensorrt as trt # noqa
@ -815,13 +818,12 @@ class Exporter:
import tensorflow as tf # noqa
check_requirements(
(
"cmake", # 'cmake' is needed to build onnxsim on aarch64 and Conda runners
"keras", # required by onnx2tf package
"tf_keras", # required by onnx2tf package
"onnx>=1.12.0",
"onnx2tf>1.17.5,<=1.22.3",
"sng4onnx>=1.0.1",
"onnxsim>=0.4.33",
"onnxslim==0.1.28",
"onnx_graphsurgeon>=0.3.26",
"tflite_support<=0.4.3" if IS_JETSON else "tflite_support", # fix ImportError 'GLIBCXX_3.4.29'
"flatbuffers>=23.5.26,<100", # update old 'flatbuffers' included inside tensorflow package