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

Install cmake with onnxsim (#13222)

Browse source
This commit is contained in:
Glenn Jocher 2024-05-30 11:36:15 +02:00 • committed by GitHub
parent 6b28bb6e48
commit fbd8fdb53e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

7
ultralytics/engine/exporter.py
View file

@ -384,9 +384,7 @@ class Exporter:
"""YOLOv8 ONNX export.""" """YOLOv8 ONNX export."""
requirements = ["onnx>=1.12.0"] requirements = ["onnx>=1.12.0"]
if self.args.simplify: if self.args.simplify:
requirements += ["onnxsim>=0.4.33", "onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime"] requirements += ["cmake", "onnxsim>=0.4.33", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
if ARM64:
check_requirements("cmake") # 'cmake' is needed to build onnxsim on aarch64
check_requirements(requirements) check_requirements(requirements)
import onnx # noqa import onnx # noqa
@ -815,10 +813,9 @@ class Exporter:
version = ">=2.0.0" version = ">=2.0.0"
check_requirements(f"tensorflow{suffix}{version}") check_requirements(f"tensorflow{suffix}{version}")
import tensorflow as tf # noqa import tensorflow as tf # noqa
if ARM64:
check_requirements("cmake") # 'cmake' is needed to build onnxsim on aarch64
check_requirements( check_requirements(
( (
"cmake", # 'cmake' is needed to build onnxsim on aarch64 and Conda runners
"keras", # required by onnx2tf package "keras", # required by onnx2tf package
"tf_keras", # required by onnx2tf package "tf_keras", # required by onnx2tf package
"onnx>=1.12.0", "onnx>=1.12.0",