diff --git a/docker/hwaccel.ml.yml b/docker/hwaccel.ml.yml index 3672014ebb..c7035047c7 100644 --- a/docker/hwaccel.ml.yml +++ b/docker/hwaccel.ml.yml @@ -17,10 +17,9 @@ services: rknn: security_opt: - systempaths=unconfined + - apparmor=unconfined devices: - /dev/dri:/dev/dri - volumes: - - /sys/kernel/debug/:/sys/kernel/debug/:ro cpu: {} diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile index 4273ca02db..ad5cba980f 100644 --- a/machine-learning/Dockerfile +++ b/machine-learning/Dockerfile @@ -106,7 +106,6 @@ RUN echo "hard core 0" >> /etc/security/limits.conf && \ COPY --from=builder /usr/src/app/.venv /usr/src/app/.venv COPY ann/ann.py /usr/src/ann/ann.py -COPY rknn/rknnpool.py /usr/src/rknn/rknnpool.py COPY start.sh log_conf.json gunicorn_conf.py ./ COPY app . diff --git a/machine-learning/app/main.py b/machine-learning/app/main.py index 0f50257a4e..4c380dc65f 100644 --- a/machine-learning/app/main.py +++ b/machine-learning/app/main.py @@ -226,9 +226,9 @@ async def load(model: InferenceModel) -> InferenceModel: except FileNotFoundError as e: if model.model_format == ModelFormat.ONNX: raise e - log.exception(e) log.warning( - f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it." + f"{model.model_format.upper()} is available, but model '{model.model_name}' does not support it.", + exc_info=e, ) model.model_format = ModelFormat.ONNX model.load() diff --git a/machine-learning/app/models/base.py b/machine-learning/app/models/base.py index 571d33de22..5ed18392ee 100644 --- a/machine-learning/app/models/base.py +++ b/machine-learning/app/models/base.py @@ -8,9 +8,8 @@ from typing import Any, ClassVar from huggingface_hub import snapshot_download import ann.ann -import rknn.rknnpool +import app.sessions.rknn as rknn from app.sessions.ort import OrtSession -from app.sessions.rknn import RknnSession from ..config import clean_name, log, settings from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType @@ -34,6 +33,7 @@ class InferenceModel(ABC): self.model_name = clean_name(model_name) self.cache_dir = Path(cache_dir) if cache_dir is not None else self._cache_dir_default self.model_format = model_format if model_format is not None else self._model_format_default + self.model_path_prefix = rknn.model_prefix if self.model_format == ModelFormat.RKNN else None if session is not None: self.session = session @@ -116,7 +116,7 @@ class InferenceModel(ABC): case ".onnx": session = OrtSession(model_path) case ".rknn": - session = RknnSession(model_path) + session = rknn.RknnSession(model_path) case _: raise ValueError(f"Unsupported model file type: {model_path.suffix}") return session @@ -127,6 +127,8 @@ class InferenceModel(ABC): @property def model_path(self) -> Path: + if self.model_path_prefix: + return self.model_dir / self.model_path_prefix / f"model.{self.model_format}" return self.model_dir / f"model.{self.model_format}" @property @@ -164,7 +166,7 @@ class InferenceModel(ABC): @property def _model_format_default(self) -> ModelFormat: - if rknn.rknnpool.is_available and settings.rknn: + if rknn.is_available: return ModelFormat.RKNN elif ann.ann.is_available and settings.ann: return ModelFormat.ARMNN diff --git a/machine-learning/app/sessions/rknn.py b/machine-learning/app/sessions/rknn/__init__.py similarity index 70% rename from machine-learning/app/sessions/rknn.py rename to machine-learning/app/sessions/rknn/__init__.py index fe490af074..a35ffcacab 100644 --- a/machine-learning/app/sessions/rknn.py +++ b/machine-learning/app/sessions/rknn/__init__.py @@ -6,15 +6,17 @@ from typing import Any, NamedTuple import numpy as np from numpy.typing import NDArray +from app.config import log, settings from app.schemas import SessionNode -from rknn.rknnpool import RknnPoolExecutor, soc_name -from ..config import log, settings +from .rknnpool import RknnPoolExecutor, is_available, soc_name + +is_available = is_available and settings.rknn +model_prefix = Path("rknpu") / soc_name if is_available else None -def runInference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: +def run_inference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw") - return outputs @@ -38,17 +40,13 @@ input_output_mapping: dict[str, dict[str, Any]] = { class RknnSession: - def __init__(self, model_path: Path | str): - self.model_path = Path(str(model_path).replace("model", soc_name)) - self.model_type = "detection" if "detection" in self.model_path.as_posix() else "recognition" + def __init__(self, model_path: Path) -> None: + self.model_type = "detection" if "detection" in model_path.parts else "recognition" self.tpe = settings.rknn_threads - log.info(f"Loading RKNN model from {self.model_path} with {self.tpe} threads.") - self.rknnpool = RknnPoolExecutor(rknnModel=self.model_path.as_posix(), tpes=self.tpe, func=runInference) - log.info(f"Loaded RKNN model from {self.model_path} with {self.tpe} threads.") - - def __del__(self) -> None: - self.rknnpool.release() + log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.") + self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference) + log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.") def get_inputs(self) -> list[SessionNode]: return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()] diff --git a/machine-learning/rknn/rknnpool.py b/machine-learning/app/sessions/rknn/rknnpool.py similarity index 58% rename from machine-learning/rknn/rknnpool.py rename to machine-learning/app/sessions/rknn/rknnpool.py index 14fa06481e..834cf39b5e 100644 --- a/machine-learning/rknn/rknnpool.py +++ b/machine-learning/app/sessions/rknn/rknnpool.py @@ -1,49 +1,55 @@ # This code is from leafqycc/rknn-multi-threaded # Following Apache License 2.0 -import os from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from queue import Queue + import numpy as np -from typing import Callable from numpy.typing import NDArray + from app.config import log supported_socs = ["rk3566", "rk3588"] -coremask_supported_socs = ["rk3576","rk3588"] +coremask_supported_socs = ["rk3576", "rk3588"] + +def get_soc(device_tree_path: Path | str) -> str | None: + try: + with Path(device_tree_path).open() as f: + device_compatible_str = f.read() + for soc in supported_socs: + if soc in device_compatible_str: + return soc + log.warning("Device is not supported for RKNN") + except OSError as e: + log.warning("Could not read /proc/device-tree/compatible. Reason: %s", e.msg) + return None + + +soc_name = None +is_available = False try: from rknnlite.api import RKNNLite - with open("/proc/device-tree/compatible") as f: - device_compatible_str = f.read() - for soc in supported_socs: - if soc in device_compatible_str: - is_available = True - soc_name = soc - break - else: - is_available = False - soc_name = None - is_available = is_available and os.path.exists("/sys/kernel/debug/rknpu/load") -except (FileNotFoundError, ImportError): + soc_name = get_soc("/proc/device-tree/compatible") + is_available = soc_name is not None +except ImportError: log.debug("RKNN is not available") - is_available = False - soc_name = None -def init_rknn(rknnModel) -> Callable: +def init_rknn(model_path: str) -> RKNNLite: if not is_available: raise RuntimeError("rknn is not available!") rknn_lite = RKNNLite() - ret = rknn_lite.load_rknn(rknnModel) + ret = rknn_lite.load_rknn(model_path) if ret != 0: raise RuntimeError("Load RKNN rknnModel failed") - + if soc_name in coremask_supported_socs: ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_AUTO) else: - ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms. + ret = rknn_lite.init_runtime() # Please do not set this parameter on other platforms. if ret != 0: raise RuntimeError("Init runtime environment failed") @@ -51,18 +57,11 @@ def init_rknn(rknnModel) -> Callable: return rknn_lite -def init_rknns(rknnModel, tpes) -> list[Callable]: - rknn_list = [] - for i in range(tpes): - rknn_list.append(init_rknn(rknnModel)) - return rknn_list - - class RknnPoolExecutor: - def __init__(self, rknnModel: str, tpes: int, func): + def __init__(self, model_path: str, tpes: int, func): self.tpes = tpes self.queue = Queue() - self.rknn_pool = init_rknns(rknnModel, tpes) + self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)] self.pool = ThreadPoolExecutor(max_workers=tpes) self.func = func self.num = 0 @@ -81,3 +80,6 @@ class RknnPoolExecutor: self.pool.shutdown() for rknn_lite in self.rknn_pool: rknn_lite.release() + + def __del__(self) -> None: + self.release() diff --git a/machine-learning/export/Dockerfile b/machine-learning/export/onnx/Dockerfile similarity index 100% rename from machine-learning/export/Dockerfile rename to machine-learning/export/onnx/Dockerfile diff --git a/machine-learning/export/conda-lock.yml b/machine-learning/export/onnx/conda-lock.yml similarity index 100% rename from machine-learning/export/conda-lock.yml rename to machine-learning/export/onnx/conda-lock.yml diff --git a/machine-learning/export/env.dev.yaml b/machine-learning/export/onnx/env.dev.yaml similarity index 100% rename from machine-learning/export/env.dev.yaml rename to machine-learning/export/onnx/env.dev.yaml diff --git a/machine-learning/export/env.yaml b/machine-learning/export/onnx/env.yaml similarity index 100% rename from machine-learning/export/env.yaml rename to machine-learning/export/onnx/env.yaml diff --git a/machine-learning/export/models/__init__.py b/machine-learning/export/onnx/models/__init__.py similarity index 100% rename from machine-learning/export/models/__init__.py rename to machine-learning/export/onnx/models/__init__.py diff --git a/machine-learning/export/models/mclip.py b/machine-learning/export/onnx/models/mclip.py similarity index 100% rename from machine-learning/export/models/mclip.py rename to machine-learning/export/onnx/models/mclip.py diff --git a/machine-learning/export/models/openclip.py b/machine-learning/export/onnx/models/openclip.py similarity index 100% rename from machine-learning/export/models/openclip.py rename to machine-learning/export/onnx/models/openclip.py diff --git a/machine-learning/export/models/optimize.py b/machine-learning/export/onnx/models/optimize.py similarity index 100% rename from machine-learning/export/models/optimize.py rename to machine-learning/export/onnx/models/optimize.py diff --git a/machine-learning/export/models/util.py b/machine-learning/export/onnx/models/util.py similarity index 100% rename from machine-learning/export/models/util.py rename to machine-learning/export/onnx/models/util.py diff --git a/machine-learning/export/run.py b/machine-learning/export/onnx/run.py similarity index 100% rename from machine-learning/export/run.py rename to machine-learning/export/onnx/run.py diff --git a/machine-learning/export/rknpu/ViT-B-32__openai/README.md b/machine-learning/export/rknpu/ViT-B-32__openai/README.md new file mode 100644 index 0000000000..eca2c7da53 --- /dev/null +++ b/machine-learning/export/rknpu/ViT-B-32__openai/README.md @@ -0,0 +1,11 @@ +--- +tags: + - immich + - clip +--- +# Model Description + +This repo contains ONNX exports for the CLIP model [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32). +It separates the visual and textual encoders into separate models for the purpose of generating image and text embeddings. + +This repo is specifically intended for use with [Immich](https://immich.app/), a self-hosted photo library. diff --git a/machine-learning/export/rknpu/build_rknn.py b/machine-learning/export/rknpu/build_rknn.py new file mode 100644 index 0000000000..507da80168 --- /dev/null +++ b/machine-learning/export/rknpu/build_rknn.py @@ -0,0 +1,69 @@ +import argparse +from pathlib import Path + +from rknn.api import RKNN + +parser = argparse.ArgumentParser("ONNX to RKNN model converter") +parser.add_argument( + "model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=Path +) +parser.add_argument("target_platform", help="target platform ex:rk3566", type=str) +args = parser.parse_args() + + +def ConvertModel(model_dir: Path, target_platform: str, dynamic_input=None): + input_path = model_dir / "model.onnx" + print(f"Converting model {input_path}") + rknn = RKNN(verbose=False) + + rknn.config( + target_platform=target_platform, + dynamic_input=dynamic_input, + enable_flash_attention=True, + # remove_reshape=True, + # model_pruning=True + ) + ret = rknn.load_onnx(model=input_path.as_posix()) + + if ret != 0: + print("Load failed!") + exit(ret) + + ret = rknn.build(do_quantization=False) + + if ret != 0: + print("Build failed!") + exit(ret) + + output_path = model_dir / "rknpu" / target_platform / "model.rknn" + output_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Exporting model {model_dir} to {output_path}") + ret = rknn.export_rknn(output_path.as_posix()) + if ret != 0: + print("Export rknn model failed!") + exit(ret) + + +textual = args.model / "textual" +visual = args.model / "visual" +detection = args.model / "detection" +recognition = args.model / "recognition" + +is_dir = [textual.is_dir(), visual.is_dir(), detection.is_dir(), recognition.is_dir()] +if not any(is_dir): + print("Unknown model") + exit(1) + +is_textual, is_visual, is_detection, is_recognition = is_dir + +if is_textual: + ConvertModel(textual, target_platform=args.target_platform) + +if is_visual: + ConvertModel(visual, target_platform=args.target_platform) + +if is_detection: + ConvertModel(detection, args.target_platform, [[[1, 3, 640, 640]]]) + +if is_recognition: + ConvertModel(recognition, args.target_platform, [[[1, 3, 112, 112]]]) diff --git a/machine-learning/rknn/export/convert.sh b/machine-learning/export/rknpu/convert.sh similarity index 100% rename from machine-learning/rknn/export/convert.sh rename to machine-learning/export/rknpu/convert.sh diff --git a/machine-learning/rknn/export/requirements.txt b/machine-learning/export/rknpu/requirements.txt similarity index 100% rename from machine-learning/rknn/export/requirements.txt rename to machine-learning/export/rknpu/requirements.txt diff --git a/machine-learning/rknn/__init__.py b/machine-learning/rknn/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/machine-learning/rknn/export/build_rknn.py b/machine-learning/rknn/export/build_rknn.py deleted file mode 100644 index 15c2da04a2..0000000000 --- a/machine-learning/rknn/export/build_rknn.py +++ /dev/null @@ -1,53 +0,0 @@ -import argparse -import os - -parser = argparse.ArgumentParser("RKNN model converting") -parser.add_argument("model", help="Directory of the model that will be exported to RKNN ex:ViT-B-32__openai.", type=str) -parser.add_argument("target_platform", help="target platform ex:rk3566", type=str) -args = parser.parse_args() - - -def ConvertModel(model_path='ViT-B-32__openai/textual/model.onnx', target_platform='rk3566', dynamic_input = None): - # E build: Repeat call the 'rknn.build' or 'rknn.hybrid_quantization_step1' is not allow! - from rknn.api import RKNN - rknn = RKNN(verbose=False) - - rknn.config(target_platform=target_platform, dynamic_input=dynamic_input) - ret = rknn.load_onnx(model=model_path) - - if ret != 0: - print("Load failed!") - exit(ret) - - ret = rknn.build(do_quantization=False) - - if ret != 0: - print("Build failed!") - exit(ret) - print(model_path.replace('model.onnx',f'{target_platform}.rknn')) - ret = rknn.export_rknn(model_path.replace('model.onnx',f'{target_platform}.rknn')) - if ret != 0: - print('Export rknn model failed!') - exit(ret) - print('done') - del rknn - del RKNN - - if not os.path.isfile(f'{model_path.replace("onnx","rknn")}'): - print(f'Dummy model not found at {model_path.replace("onnx","rknn")}, creating one') - with open(f'{model_path.replace("onnx","rknn")}', 'w'): - pass - - -if os.path.isdir(f'{args.model}/textual') and os.path.isdir(f'{args.model}/visual'): # is a clip model - print('Converting Clip model.') - ConvertModel(model_path=f'{args.model}/textual/model.onnx', target_platform=args.target_platform) - ConvertModel(model_path=f'{args.model}/visual/model.onnx', target_platform=args.target_platform) - -elif os.path.isdir(f'{args.model}/detection') and os.path.isdir(f'{args.model}/recognition'): # is a facial model - print('Converting facial model.') - ConvertModel(f'{args.model}/detection/model.onnx', args.target_platform, [[[1, 3, 640, 640]]]) - ConvertModel(f'{args.model}/recognition/model.onnx', args.target_platform, [[[1, 3, 112, 112]]]) - -else: - print('Unknown model.')