mirror of
https://github.com/immich-app/immich.git
synced 2026-06-14 21:11:24 +08:00
41580696c7
Some checks are pending
CodeQL / Analyze (javascript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
Docker / Build and Push (., cpu, server/Dockerfile, immich-server, linux/amd64,linux/arm64) (push) Waiting to run
Docker / Build and Push (machine-learning, armnn, machine-learning/Dockerfile, immich-machine-learning, linux/arm64, -armnn) (push) Waiting to run
Docker / Build and Push (machine-learning, cpu, machine-learning/Dockerfile, immich-machine-learning, linux/amd64,linux/arm64) (push) Waiting to run
Docker / Build and Push (machine-learning, cuda, machine-learning/Dockerfile, immich-machine-learning, linux/amd64, -cuda) (push) Waiting to run
Docker / Build and Push (machine-learning, openvino, machine-learning/Dockerfile, immich-machine-learning, linux/amd64, -openvino) (push) Waiting to run
Static Code Analysis / Run Dart Code Analysis (push) Waiting to run
Test / Server (push) Waiting to run
Test / CLI (push) Waiting to run
Test / CLI (Windows) (push) Waiting to run
Test / Web (push) Waiting to run
Test / End-to-End Tests (push) Waiting to run
Test / Mobile (push) Waiting to run
Test / Machine Learning (push) Waiting to run
Test / ShellCheck (push) Waiting to run
Test / OpenAPI Clients (push) Waiting to run
Test / TypeORM Checks (push) Waiting to run
* update export code * add uuid glob, sort model names * add new models to ml, sort names * add new models to server, sort by dims and name * typo in name * update export dependencies * onnx save function * format
50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from pathlib import Path
|
|
|
|
import onnx
|
|
import onnxruntime as ort
|
|
import onnxsim
|
|
|
|
|
|
def save_onnx(model: onnx.ModelProto, output_path: Path | str) -> None:
|
|
try:
|
|
onnx.save(model, output_path)
|
|
except ValueError as e:
|
|
if "The proto size is larger than the 2 GB limit." in str(e):
|
|
onnx.save(model, output_path, save_as_external_data=True, size_threshold=1_000_000)
|
|
else:
|
|
raise e
|
|
|
|
|
|
def optimize_onnxsim(model_path: Path | str, output_path: Path | str) -> None:
|
|
model_path = Path(model_path)
|
|
output_path = Path(output_path)
|
|
model = onnx.load(model_path.as_posix())
|
|
model, check = onnxsim.simplify(model)
|
|
assert check, "Simplified ONNX model could not be validated"
|
|
for file in model_path.parent.iterdir():
|
|
if file.name.startswith("Constant") or "onnx" in file.name or file.suffix == ".weight":
|
|
file.unlink()
|
|
save_onnx(model, output_path)
|
|
|
|
|
|
def optimize_ort(
|
|
model_path: Path | str,
|
|
output_path: Path | str,
|
|
level: ort.GraphOptimizationLevel = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC,
|
|
) -> None:
|
|
model_path = Path(model_path)
|
|
output_path = Path(output_path)
|
|
|
|
sess_options = ort.SessionOptions()
|
|
sess_options.graph_optimization_level = level
|
|
sess_options.optimized_model_filepath = output_path.as_posix()
|
|
|
|
ort.InferenceSession(model_path.as_posix(), providers=["CPUExecutionProvider"], sess_options=sess_options)
|
|
|
|
|
|
def optimize(model_path: Path | str) -> None:
|
|
model_path = Path(model_path)
|
|
|
|
optimize_ort(model_path, model_path)
|
|
optimize_onnxsim(model_path, model_path)
|