mlonmcu 0.3.0.dev0__3-none-any.whl

mlonmcu/models/lookup.py

@@ -0,0 +1,345 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from pathlib import Path
+import os
+import glob
+from itertools import product
+import yaml
+
+from .model import Model, ModelFormats
+from .group import ModelGroup
+
+from mlonmcu.logging import get_logger
+
+logger = get_logger()
+
+
+def get_model_directories(context):
+    dirs = context.environment.paths["models"]
+    dirs = [d.path for d in dirs]
+    return dirs
+
+
+def find_metadata(directory, model_name=None):
+    possible_basenames = ["model", "metadata", "definition"]
+    possible_extensions = ["yaml", "yml"]
+    directory = Path(directory)
+    dirname = directory.name
+    if dirname not in possible_basenames:
+        possible_basenames.insert(0, dirname)
+    if model_name and model_name not in possible_basenames:
+        possible_basenames.insert(0, model_name)
+
+    for combination in product(possible_basenames, possible_extensions):
+        filename = f"{combination[0]}.{combination[1]}"
+        fullpath = directory / filename
+        if fullpath.is_file():
+            return fullpath
+            # logger.debug("Found match. Ignoring other files")
+    return None
+
+
+def list_models(directory, depth=1, formats=None, config=None):  # TODO: get config from environment or cmdline!
+    config = config if config is not None else {}
+    formats = formats if formats else [ModelFormats.TFLITE]
+    assert len(formats) > 0, "No formats provided for model lookup"
+    models = []
+    for fmt in formats:
+        if depth != 1:
+            raise NotImplementedError  # TODO: implement for arm ml zoo
+            # define all allowed extensions + search recusively (with limit?)
+            # list(Path(".").rglob(f"*.{ext}")) for ext in allowed_ext
+        if not os.path.isdir(directory):
+            logger.debug("Not a directory: %s", str(directory))
+            return []
+        subdirs = [Path(directory) / o for o in os.listdir(directory) if os.path.isdir(os.path.join(directory, o))]
+        for subdir in subdirs:
+            dirname = subdir.name
+            if dirname.startswith("."):
+                # logger.debug("Skipping hidden directory: %s", str(dirname))
+                continue
+            exts = fmt.extensions
+            for ext in exts:
+                main_model = subdir / f"{dirname}.{ext}"
+                if os.path.exists(main_model):
+                    main_model = f"{dirname}/{dirname}"
+                else:
+                    main_model = None
+                submodels = []
+                for filename in glob.glob(str(subdir / f"*.{ext}")):
+                    basename = "".join(Path(filename).name.split(".")[:-1])
+                    submodels.append(f"{dirname}/{basename}")
+
+                if len(submodels) == 1:
+                    main_model = submodels[0]
+
+                if main_model:
+                    submodels.remove(main_model)
+
+                    main_base = main_model.split("/")[-1]
+                    main_config = {}
+                    main_metadata_path = find_metadata(Path(directory) / dirname, model_name=main_base)
+                    if main_metadata_path:
+                        main_config[f"{main_base}.metadata_path"] = main_metadata_path
+                    main_config.update(config)
+
+                    models.append(
+                        Model(
+                            main_base,
+                            [Path(directory) / f"{main_model}.{ext}"],
+                            config=main_config,
+                            alt=main_model,
+                            formats=[fmt],
+                        )
+                    )
+
+                for submodel in submodels:
+                    sub_base = submodel.split("/")[-1]
+                    submodel_config = {}
+                    submodel_metadata_path = find_metadata(Path(directory) / dirname, model_name=sub_base)
+                    if submodel_metadata_path:
+                        submodel_config[f"{sub_base}.metadata_path"] = submodel_metadata_path
+                    submodel_config.update(config)
+                    models.append(
+                        Model(
+                            submodel,
+                            [Path(directory) / f"{submodel}.{ext}"],
+                            config=submodel_config,
+                            formats=[fmt],
+                        )
+                    )
+
+    return models
+
+
+def list_modelgroups(directory):
+    if not os.path.isdir(directory):
+        logger.debug("Not a directory: %s", str(directory))
+        return []
+    groups = []
+    directory = Path(directory)
+    possible_basenames = ["modelgroups", "groups"]
+    possible_extensions = ["yaml", "yml"]
+    for combination in product(possible_basenames, possible_extensions):
+        filename = f"{combination[0]}.{combination[1]}"
+        fullpath = directory / filename
+        if fullpath.is_file():
+            logger.debug("Found match. Ignoring other files")
+            with open(fullpath, "r") as yamlfile:
+                try:
+                    content = yaml.safe_load(yamlfile)
+                    for groupname, groupmodels in content.items():
+                        assert isinstance(groupmodels, list), "Modelgroups should be defined as a YAML list"
+                        modelgroup = ModelGroup(groupname, groupmodels)
+                        groups.append(modelgroup)
+                except yaml.YAMLError as err:
+                    raise RuntimeError("Could not open YAML file") from err
+            break
+    return groups
+
+
+def lookup_models_and_groups(directories, formats):
+    all_models = []
+    all_groups = []
+    duplicates = {}
+    group_duplicates = {}
+    for directory in directories:
+        models = list_models(directory, formats=formats)
+        if len(all_models) == 0:
+            all_models = models
+        else:
+            all_model_names = [m.name for m in all_models]
+            for model in models:
+                name = model.name
+                if name in all_model_names:
+                    if name in duplicates:
+                        duplicates[name] += 1
+                    else:
+                        duplicates[name] = 1
+                else:
+                    all_models.append(model)
+        groups = list_modelgroups(directory)
+        if len(all_groups) == 0:
+            all_groups = groups
+        else:
+            all_group_names = [g.name for g in all_groups]
+            for group in groups:
+                name = group.name
+                if name in all_group_names:
+                    if name in group_duplicates:
+                        group_duplicates[name] += 1
+                    else:
+                        group_duplicates[name] = 1
+                else:
+                    all_groups.append(group)
+
+    all_models = sorted(all_models, key=lambda x: x.name)
+
+    return all_models, all_groups, duplicates, group_duplicates
+
+
+def print_paths(directories):
+    print("Paths:")
+    for directory in directories:
+        exists = os.path.exists(directory)
+        print("    " + str(directory), end="\n" if exists else " (skipped)\n")
+        # if not exists:
+        #     directories.remove(directory)
+    print()
+
+
+def print_models(models, duplicates=[], detailed=False):
+    print("Models:")
+    for model in models:
+        name = model.name
+        path = model.paths[0]
+        meta = "available" if model.metadata is not None else "not available"
+        print("    " + name, end="")
+        if name in duplicates:
+            num = duplicates[name]
+            print(f" ({num} duplicates)")
+        else:
+            print()
+        if detailed:
+            print(f"        Path: {path}")
+            print(f"        Metadata: {meta}")
+    print()
+
+
+def print_groups(groups, all_models=[], duplicates=[], detailed=False):
+    print("Groups:")
+    for group in groups:
+        name = group.name
+        models = group.models
+        all_model_names = [m.name for m in all_models]
+        for model in models:
+            if model in all_model_names:
+                # Groupname conflicts with modelname
+                group.name = "~" + group.name
+        size = len(models)
+        print(f"    {name} [{size} models]", end="")
+        if name in duplicates:
+            num = duplicates[name]
+            print(f" ({num} duplicates)")
+        else:
+            print()
+
+        if detailed:
+            groupmodels = " ".join(models)
+            print(f"        Models: {groupmodels}")
+            print()
+
+
+def print_summary(context, detailed=False):
+    # TODO: get from context!
+    formats = [ModelFormats.TFLITE]
+
+    directories = get_model_directories(context)
+
+    models, groups, duplicates, group_duplicates = lookup_models_and_groups(directories, formats)
+
+    print("Models Summary\n")
+    print_paths(directories)
+    print_models(models, duplicates=duplicates, detailed=detailed)
+    print_groups(groups, duplicates=group_duplicates, all_models=models, detailed=detailed)
+
+
+def lookup_models(names, frontends=None, context=None):
+    if frontends is None:
+        assert context is not None
+        # TODO: Get defaults frontends from environment (with no config/features)
+        raise NotImplementedError
+        # frontends = ?
+    # allowed_ext = [frontend.fmt.extension for frontend in frontends]
+    allowed_fmts = list(dict.fromkeys(sum([frontend.input_formats for frontend in frontends], [])))  # Remove duplicates
+    allowed_exts = [fmt.extension for fmt in allowed_fmts]  # There should nit be duplicates
+
+    if context:
+        directories = get_model_directories(context)
+        models, _, _, _ = lookup_models_and_groups(directories, allowed_fmts)
+    else:
+        models = []
+    model_names = [model.name for model in models]
+
+    hints = []
+    for name in names:
+        filepath = Path(name)
+        real_name = filepath.stem
+        ext = filepath.suffix[1:]
+        if len(ext) > 0 and filepath.is_file():  # Explicit file
+            assert (
+                ext in allowed_exts
+            ), f"Unsupported file extension for model which was explicitly passed by path: {ext}"
+            paths = [filepath]
+            config = {}
+            metadata_path = find_metadata(filepath.parent.resolve(), model_name=real_name)
+            if metadata_path:
+                config[f"{real_name}.metadata_path"] = metadata_path
+            hint = Model(real_name, paths, config=config, formats=[ModelFormats.from_extension(ext)])
+            # TODO: look for metadata
+            hints.append(hint)
+        else:
+            assert context is not None, "Context is required for passing models by name"
+            assert len(model_names) > 0, "List of available models is empty"
+            assert name in model_names, f"Could not find a model matching the name: {name}"
+            index = model_names.index(name)
+            hint = models[index]
+            hints.append(hint)
+
+    return hints
+
+
+def apply_modelgroups(models, context=None):
+    assert context is not None
+    directories = get_model_directories(context)
+
+    groups = {}
+    for directory in directories:
+        for group in list_modelgroups(directory):
+            if group.name in groups and groups[group.name] != group.models:
+                raise RuntimeError(
+                    f"The model group '{group.name}' has conflicting definitions. Used the following directories:"
+                    f" {directories}"
+                )
+            groups[group.name] = group.models
+
+    out_models = []
+    for model in models:
+        out_models.extend(groups.get(model, [model]))
+    return list(dict.fromkeys(out_models))  # Drop duplicates
+
+
+def map_frontend_to_model(model, frontends, backend=None):
+    model_fmts = model.formats
+    backend_fmts = backend.get_supported_formats() if backend else []
+    for frontend in frontends:
+        ins = frontend.input_formats
+        outs = frontend.output_formats
+        assert len(ins) > 0
+        # TODO: instead of picking the first frontend which at least one match, determine the highest overlap
+        if any(in_fmt in model_fmts for in_fmt in ins):
+            if len(backend_fmts) == 0 or any(out_fmt in backend_fmts for out_fmt in outs):
+                new_model = model  # TODO: deepcopy?
+                for i, fmt in enumerate(model_fmts):
+                    if fmt not in ins:  # Only keep those formqts which are supported by the chosen frontend
+                        path = model.paths[model_fmts.index(fmt)]
+                        new_model.formats.remove(fmt)
+                        new_model.paths.remove(path)
+                return new_model, frontend
+    raise RuntimeError(f"Unable to find a suitable frontend for model '{model.name}'")

mlonmcu/models/metadata.py

@@ -0,0 +1,133 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import yaml
+
+
+def parse_metadata(path):
+    content = {}
+    with open(path, "r") as yamlfile:
+        try:
+            content = yaml.safe_load(yamlfile)
+        except yaml.YAMLError as err:  # pragma: no cover
+            raise RuntimeError(f"Could not open YAML file: {path}") from err
+    return content
+
+
+# class ModelOperator:
+#     def __init__(self, name, custom=False, registration=None):
+#         self.name = name
+#         self.custom = custom
+#         self.registration = registration
+#
+#
+# class ModelMetadata:
+#     def __init__(
+#         self,
+#         author="unknown",
+#         description="",
+#         created_at=None,
+#         references=None,
+#         comment=None,
+#         backend_options_map=None,
+#     ):
+#         self.author = author if author is not None else "unknown"
+#         self.description = description if description is not None else ""
+#         self.created_at = created_at
+#         self.refrences = references if references is not None else []
+#         self.comment = comment
+#         self.backend_options_map = (
+#             backend_options_map if backend_options_map is not None else {}
+#         )
+#
+#     def __repr__(self):
+#         return "ModelMetadata(" + str(vars(self)) + ")"
+#
+#
+# class TfLitelMetadata(ModelMetadata):
+#     def __init__(
+#         self,
+#         author="unknown",
+#         description="",
+#         created_at=None,
+#         references=None,
+#         comment=None,
+#         backend_options_map=None,
+#         operators=[],
+#     ):
+#         super().__init__(
+#             author=author,
+#             description=description,
+#             created_at=created_at,
+#             references=references,
+#             comment=comment,
+#             backend_options_map=backend_options_map,
+#             operators=operators,
+#         )
+#         self.operators = operators
+#
+#
+# def parse_metadata(path):
+#     with open(path, "r") as yamlfile:
+#         try:
+#             content = yaml.safe_load(yamlfile)
+#             if not content:
+#                 # file empty
+#                 return ModelMetadata()
+#             if "author" in content:
+#                 author = content["author"]
+#             else:
+#                 author = None
+#             if "description" in content:
+#                 description = content["description"]
+#             else:
+#                 description = None
+#             if "created_at" in content:
+#                 created_at = content["created_at"]
+#             else:
+#                 created_at = None
+#             if "references" in content:
+#                 references = content["references"]
+#                 assert isinstance(references, list)
+#             else:
+#                 references = []
+#             if "comment" in content:
+#                 comment = content["comment"]
+#             else:
+#                 comment = None
+#             backend_options_map = {}
+#             if "backends" in content:
+#                 backends = content["backends"]
+#                 for backend in backends:
+#                     backend_options = parse_model_options_for_backend(
+#                         backend, backends[backend]
+#                     )
+#                     backend_options_map[backend] = backend_options
+#             metadata = ModelMetadata(
+#                 author=author,
+#                 description=description,
+#                 created_at=created_at,
+#                 references=references,
+#                 comment=comment,
+#                 backend_options_map=backend_options_map,
+#             )
+#             return metadata
+#         except yaml.YAMLError as err:
+#             raise RuntimeError(f"Could not open YAML file: {path}") from err
+#
+#     return None

mlonmcu/models/model.py

@@ -0,0 +1,104 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from enum import Enum
+from pathlib import Path
+from collections import namedtuple
+
+from mlonmcu.config import filter_config
+
+from .metadata import parse_metadata
+
+ModelFormat = namedtuple("ModelFormat", ["value", "extensions"])
+
+
+class ModelFormats(Enum):
+    @property
+    def extensions(self):
+        return self.value.extensions
+
+    @property
+    def extension(self):
+        return self.value.extensions[0]
+
+    @classmethod
+    def from_extension(cls, ext):
+        for fmt in cls:
+            if ext in fmt.extensions:
+                return fmt
+        return None
+
+    NONE = ModelFormat(0, [])
+    TFLITE = ModelFormat(1, ["tflite"])
+    PACKED = ModelFormat(2, ["tflm"])
+    IPYNB = ModelFormat(3, ["ipynb"])
+    ONNX = ModelFormat(4, ["onnx"])
+    RELAY = ModelFormat(5, ["relay"])
+    PB = ModelFormat(6, ["pb"])
+    PADDLE = ModelFormat(7, ["pdmodel"])
+
+
+def parse_metadata_from_path(path):
+    if Path(path).is_file():
+        metadata = parse_metadata(path)
+        return metadata
+    return None
+
+
+class Model:
+    DEFAULTS = {
+        "metadata_path": "definition.yml",
+        "support_path": "support",
+        "inputs_path": "input",
+        "outputs_path": "output",
+    }
+
+    def __init__(self, name, paths, config=None, alt=None, formats=ModelFormats.TFLITE):
+        self.name = name
+        self.paths = paths
+        if not isinstance(self.paths, list):
+            self.paths = [self.path]
+        self.alt = alt
+        self.formats = formats
+        if not isinstance(self.formats, list):
+            self.formats = [formats]
+        self.config = filter_config(config if config is not None else {}, self.name, self.DEFAULTS, [], [])
+        self.metadata = parse_metadata_from_path(self.metadata_path)
+
+    @property
+    def metadata_path(self):
+        return self.config["metadata_path"]
+
+    @property
+    def support_path(self):
+        return self.config["support_path"]
+
+    @property
+    def inputs_path(self):
+        # TODO: fall back to metadata
+        return self.config["inputs_path"]
+
+    @property
+    def outputs_path(self):
+        # TODO: fall back to metadata
+        return self.config["outputs_path"]
+
+    def __repr__(self):
+        if self.alt:
+            return f"Model({self.name},alt={self.alt})"
+        return f"Model({self.name})"

mlonmcu/models/options.py

@@ -0,0 +1,64 @@
+#
+# Copyright (c) 2022 TUM Department of Electrical and Computer Engineering.
+#
+# This file is part of MLonMCU.
+# See https://github.com/tum-ei-eda/mlonmcu.git for further info.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+class BackendModelOptions:
+    def __init__(self, backend, supported=True, options={}):
+        self.backend = backend
+        self.supported = supported
+        self.options = options
+
+
+class TFLMIModelOptions(BackendModelOptions):
+    def __init__(
+        self,
+        backend,
+        supported=True,
+        arena_size=None,
+        builtin_ops=None,
+        custom_ops=None,
+    ):
+        super().__init__(backend, supported=supported)
+        self.arena_size = arena_size
+        self.builtin_ops = builtin_ops
+        self.custom_ops = custom_ops
+
+
+class TVMRTModelOptions(BackendModelOptions):
+    def __init__(self, backend, supported=True, arena_size=None):
+        super().__init__(backend, supported=supported)
+        self.arena_size = arena_size
+
+
+def parse_model_options_for_backend(backend, options):
+    backend_types = {
+        "tflmi": TFLMIModelOptions,
+        "tvmrt": TVMRTModelOptions,
+    }
+    if backend in backend_types:
+        backend_type = backend_types[backend]
+    else:
+        backend_type = BackendModelOptions
+
+    backend_options = backend_type(backend)
+
+    for key, value in options.items():
+        setattr(backend_options, key, value)
+
+    return backend_options