sparsezoo.deployment_package

MANIFEST.in

@@ -1 +1,2 @@
 include LICENSE
+include src/sparsezoo/deployment_package/docker/*

setup.py

@@ -92,6 +92,7 @@ def _setup_entry_points() -> Dict:
         "console_scripts": [
             "sparsezoo=sparsezoo.main:main",
             "sparsezoo.download=sparsezoo.download_main:main",
+            "sparsezoo.deployment_package=sparsezoo.deployment_package.cli:main",
         ]
     }
 

src/sparsezoo/__init__.py

@@ -13,10 +13,13 @@
 # limitations under the License.
 
 # flake8: noqa
+# isort: skip_file
 
 from .inference import *
 from .model import *
 from .objects import *
 from .search import *
 from .utils import *
 from .validation import *
+from . import deployment_package as deployment_package_module
+from .deployment_package import *

src/sparsezoo/deployment_package/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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.
+# flake8: noqa
+
+from .main import *

src/sparsezoo/deployment_package/cli.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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.
+
+# flake8:noqa
+"""
+Usage: sparsezoo.deployment_package [OPTIONS] [DIRECTORY]
+
+  Utility to fetch a deployment directory for a task based on a optimizing-
+  metric
+
+  Example for using sparsezoo.deployment_package:
+
+       1) `sparsezoo.deployment_package --task image_classification -m accuracy`
+
+       2) `sparsezoo.deployment_package --task ic --optimizing_metric accuracy
+       --optimizing_metric compression --target VNNI`
+
+Options:
+  --version                       Show the version and exit.
+  --task [ic|image-classification|image_classification|classification|od|
+  object-detection|object_detection|detection|segmentation|qa|question-answering|
+  question_answering|text-classification|text_classification|glue|sentiment|
+  sentiment_analysis|sentiment-analysis|token-classification|token_classification|
+  ner|named-entity-recognition|named_entity_recognition]
+                                  The task to find model for, must be
+                                  specified if `--dataset` not provided
+  --dataset [imagenette|imagenet|coco|squad|mnli|qqp|sst2|conll2003]
+                                  The public dataset used to train this model,
+                                  must be specified if `--task` not provided
+  --optimizing-metric, --optimizing_metric TEXT
+                                  The criterion to search model for, multiple
+                                  metrics can be specified as comma
+                                  separated values, supported metrics are
+                                  ['accuracy', 'f1', 'recall', 'mAP',
+                                  'latency', 'throughput', 'compression',
+                                  'file_size', 'memory_usage']  [default:
+                                  accuracy]
+  --target [VNNI|DEFAULT]         Deployment target scenario (ie 'VNNI' for
+                                  VNNI capable CPUs)  [default: DEFAULT]
+  --help                          Show this message and exit.  [default:
+                                  False]
+##########
+Examples:
+    1) Fetch the smallest Image Classification Model trained on imagenette
+        sparsezoo.deployment_package --dataset imagenette --optimizing_metric compression
+    2) Fetch the most accurate Image Classification  Model trained on imagenette
+        sparsezoo.deployment_package --dataset imagenette --optimizing_metric accuracy
+    3) Fetch the most performant Question Answering model trained on squad
+        sparsezoo.deployment_package --task qa --dataset squad --optimizing_metric latency
+    4) Fetch the smallest most performant Question Answering model trained on squad
+        sparsezoo.deployment_package --task qa --dataset squad \
+            --optimizing_metric "compression, accuracy"
+"""
+import logging
+from typing import Any, Mapping, Optional
+
+import click
+from sparsezoo import Model, deployment_package
+from sparsezoo.deployment_package.docker.helpers import DEPLOYMENT_DOCKER_PATH
+from sparsezoo.utils import (
+    DATASETS,
+    DEFAULT_DEPLOYMENT_SCENARIO,
+    DEFAULT_OPTIMIZING_METRIC,
+    DEPLOYMENT_SCENARIOS,
+    METRICS,
+    TASKS_WITH_ALIASES,
+)
+from sparsezoo.version import __version__
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def _csv_callback(ctx, self, value):
+    """
+    A click callback function to parse a comma separated string with metrics
+    into a list
+    """
+    current_metrics = []
+    for metric in value.split(","):
+        metric_ = metric.lower().strip()
+        if metric_ not in METRICS:
+            raise ValueError(f"Specified metric {metric_} is not supported")
+        current_metrics.append(metric_)
+    return current_metrics
+
+
+def _download_deployment_directory(stub: str, destination: Optional[str] = None) -> str:
+    model = Model(stub)
+    model.deployment.download(destination_path=destination)
+    return model.deployment.path
+
+
+def _get_template(results: Mapping[str, Any], destination: Optional[str] = None):
+    stub = results.get("stub")
+
+    if not stub:
+        return "No relevant models found for specified metrics"
+
+    stub_info = f"""
+        Relevant Stub: {stub}
+    """
+    metrics = results.get("metrics")
+    metrics_info = (
+        f"""
+        Model Metrics: {metrics}
+        """
+        if metrics
+        else ""
+    )
+    dockerfile = DEPLOYMENT_DOCKER_PATH
+    dockerfile_directory = DEPLOYMENT_DOCKER_PATH.parent
+    deployment_path = _download_deployment_directory(stub, destination=destination)
+    download_instructions = f"""
+        Use the dockerfile in {dockerfile} to build deepsparse
+        image and run the `deepsparse.server` container. 
+
+        Run the following command inside `{dockerfile_directory}`
+        directory (Note: replace <TASK-NAME> with appropriate task):
+
+        ```bash        
+        docker build -t deepsparse_docker . && docker run -it \\
+        -v {deployment_path}:/home/deployment  deepsparse_docker \\
+         deepsparse.server --task <TASK-NAME> --model_path /home/deployment
+        ```
+    """
+    return "".join((stub_info, metrics_info, download_instructions))
+
+
+@click.command(
+    context_settings=dict(
+        token_normalize_func=lambda x: x.replace("-", "_"),
+        show_default=True,
+    )
+)
+@click.version_option(version=__version__)
+@click.argument(
+    "directory",
+    type=str,
+    default="",  # defaulting to `None` throws a missing argument Error
+)
+@click.option(
+    "--task",
+    type=click.Choice(TASKS_WITH_ALIASES, case_sensitive=False),
+    help="The task to find model for, must be specified if `--dataset` not provided",
+)
+@click.option(
+    "--dataset",
+    type=click.Choice(DATASETS, case_sensitive=False),
+    help="The public dataset used to train this model, must be specified if "
+    "`--task` not provided",
+)
+@click.option(
+    "--optimizing-metric",
+    "-m",
+    default=DEFAULT_OPTIMIZING_METRIC,
+    type=str,
+    help="The criterion to search model for, multiple metrics can be specified "
+    f"as comma separated values, supported metrics are {METRICS}",
+    callback=_csv_callback,
+)
+@click.option(
+    "--target",
+    type=click.Choice(DEPLOYMENT_SCENARIOS, case_sensitive=False),
+    default=DEFAULT_DEPLOYMENT_SCENARIO,
+    help="Deployment target scenario (ie 'VNNI' for VNNI capable CPUs)",
+    show_default=True,
+)
+def main(**kwargs):
+    r"""
+    Utility to fetch a deployment directory for a task based on specified
+    optimizing-metric
+
+    Example for using sparsezoo.deployment_package:
+
+         1) `sparsezoo.deployment_package --task image_classification \
+                --optimizing_metric accuracy
+
+         2) `sparsezoo.deployment_package --task ic \
+                --optimizing_metric "accuracy, compression" \
+                --target VNNI`
+    """
+    if not (kwargs.get("task") or kwargs.get("dataset")):
+        raise ValueError("At-least one of the `task` or `dataset` must be specified")
+    _LOGGER.debug(f"{kwargs}")
+    results = deployment_package(**kwargs)
+
+    print(_get_template(results=results, destination=kwargs.get("directory")))
+
+
+if __name__ == "__main__":
+    main()

src/sparsezoo/deployment_package/docker/Dockerfile

@@ -0,0 +1,25 @@
+# Setup the base image
+FROM python:3.8-slim-bullseye
+
+# Install git
+RUN : \
+    && apt-get update \
+    && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends git \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# Activate venv
+RUN python3.8 -m venv /venv
+ENV PATH="venv/bin:$PATH"
+
+RUN pip3 install --upgrade setuptools wheel pip3
+
+# Setup DeepSparse
+
+ARG GIT_CHECKOUT
+# if $GIT_CHECKOUT is not specified - just install from pypi
+RUN if [ -z "${GIT_CHECKOUT}" ] ; then pip3 install --no-cache-dir --upgrade deepsparse[server] ; fi
+
+# if $GIT_CHECKOUT is specified - clone, checkout $GIT_CHECKOUT, and install with -e
+RUN if [ -n "${GIT_CHECKOUT}" ] ; then git clone https://github.com/neuralmagic/deepsparse.git --depth 1 -b $GIT_CHECKOUT; fi
+RUN if [ -n "${GIT_CHECKOUT}" ] ; then pip3 install --no-cache-dir --upgrade -e "./deepsparse[server]" ; fi

src/sparsezoo/deployment_package/docker/helpers.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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.
+
+__all__ = ["DEPLOYMENT_DOCKER_PATH"]
+
+from pathlib import Path
+
+
+def _get_dockerfile_path():
+    top_level_dir = Path(__file__).parent
+    dockerfile_path = top_level_dir / "Dockerfile"
+    return dockerfile_path.absolute()
+
+
+DEPLOYMENT_DOCKER_PATH: Path = _get_dockerfile_path()

src/sparsezoo/deployment_package/main.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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 doctest
+
+
+__all__ = [
+    "deployment_package",
+]
+
+from typing import Any, Iterable, Mapping, Optional, Union
+
+from sparsezoo.deployment_package.utils.utils import recommend_stub
+
+
+def deployment_package(
+    task: Optional[str] = None,
+    dataset: Optional[str] = None,
+    scenario: Optional[str] = None,
+    optimizing_metric: Optional[Union[Iterable[str], str]] = None,
+    **kwargs,
+) -> Mapping[str, Any]:
+    """
+    A function that returns appropriate SparseZoo stub or deployment directory given
+    the task or dataset, optimizing criterions and a deployment scenario
+
+    :param task: str A supported task
+    :param dataset: str The public dataset this model was trained for
+    :param scenario: Optional[str] `VNNI` or `vnni for a VNNI compatible machine
+    :param optimizing_metric: an optional string or list of strings
+        representing different metrics to prioritize for when searching for models
+    :return: A dict type object with the relevant stub and model metrics
+    """
+    optimizing_metric = (
+        [optimizing_metric] if isinstance(optimizing_metric, str) else optimizing_metric
+    )
+
+    stub, metrics = recommend_stub(
+        task=task,
+        dataset=dataset,
+        optimizing_metrics=optimizing_metric,
+        scenario=scenario,
+    )
+
+    return {
+        "stub": stub,
+        "metrics": metrics,
+    }
+
+
+if __name__ == "__main__":
+    doctest.testmod()

src/sparsezoo/deployment_package/utils/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# 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.
+# flake8: noqa
+
+from .extractors import *
+from .utils import *