Deepsparse framework and sparsification implementation for phase 2 (#169)

* Deepsparse framework and sparsification implementation for phase 2

* add nightly as alternative name for deepsparse

Author: markurtz

Makefile

@@ -9,8 +9,11 @@ MDCHECKFILES := CODE_OF_CONDUCT.md CONTRIBUTING.md DEVELOPING.md README.md
 SPARSEZOO_TEST_MODE := "true"
 
 BUILD_ARGS :=  # set nightly to build nightly release
-TARGETS := ""  # targets for running pytests: keras,onnx,pytorch,pytorch_models,pytorch_datasets,tensorflow_v1,tensorflow_v1_models,tensorflow_v1_datasets
+TARGETS := ""  # targets for running pytests: deepsparse,keras,onnx,pytorch,pytorch_models,pytorch_datasets,tensorflow_v1,tensorflow_v1_models,tensorflow_v1_datasets
 PYTEST_ARGS := ""
+ifneq ($(findstring deepsparse,$(TARGETS)),deepsparse)
+    PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/deepsparse
+endif
 ifneq ($(findstring keras,$(TARGETS)),keras)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparseml/keras
 endif

src/sparseml/deepsparse/__init__.py

@@ -0,0 +1,23 @@
+# 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.
+
+"""
+Functionality for working with and sparsifying Models in the DeepSparse framework
+"""
+
+# flake8: noqa
+
+from .base import *
+from .framework import detect_framework, framework_info, is_supported
+from .sparsification import sparsification_info

src/sparseml/deepsparse/base.py

@@ -0,0 +1,105 @@
+# 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 functools
+from typing import Optional
+
+from sparseml.base import check_version
+
+
+try:
+    import deepsparse
+
+    deepsparse_err = None
+except Exception as err:
+    deepsparse = object()  # TODO: populate with fake object for necessary imports
+    deepsparse_err = err
+
+
+__all__ = [
+    "deepsparse",
+    "deepsparse_err",
+    "check_deepsparse_install",
+    "require_deepsparse",
+]
+
+
+def check_deepsparse_install(
+    min_version: Optional[str] = None,
+    max_version: Optional[str] = None,
+    raise_on_error: bool = True,
+) -> bool:
+    """
+    Check that the deepsparse package is installed.
+    If raise_on_error, will raise an ImportError if it is not installed or
+    the required version range, if set, is not installed.
+    If not raise_on_error, will return True if installed with required version
+    and False otherwise.
+
+    :param min_version: The minimum version for deepsparse that it must be greater than
+        or equal to, if unset will require no minimum version
+    :type min_version: str
+    :param max_version: The maximum version for deepsparse that it must be less than
+        or equal to, if unset will require no maximum version.
+    :type max_version: str
+    :param raise_on_error: True to raise any issues such as not installed,
+        minimum version, or maximum version as ImportError. False to return the result.
+    :type raise_on_error: bool
+    :return: If raise_on_error, will return False if deepsparse is not installed
+        or the version is outside the accepted bounds and True if everything is correct.
+    :rtype: bool
+    """
+    if deepsparse_err is not None:
+        if raise_on_error:
+            raise deepsparse_err
+        return False
+
+    return check_version(
+        "deepsparse",
+        min_version,
+        max_version,
+        raise_on_error,
+        alternate_package_names=["deepsparse-nightly"],
+    )
+
+
+def require_deepsparse(
+    min_version: Optional[str] = None, max_version: Optional[str] = None
+):
+    """
+    Decorator function to require use of deepsparse.
+    Will check that deepsparse package is installed and within the bounding
+    ranges of min_version and max_version if they are set before calling
+    the wrapped function.
+    See :func:`check_deepsparse_install` for more info.
+
+    param min_version: The minimum version for deepsparse that it must be greater than
+        or equal to, if unset will require no minimum version
+    :type min_version: str
+    :param max_version: The maximum version for deepsparse that it must be less than
+        or equal to, if unset will require no maximum version.
+    :type max_version: str
+    """
+
+    def _decorator(func):
+        @functools.wraps(func)
+        def _wrapper(*args, **kwargs):
+            check_deepsparse_install(min_version, max_version)
+
+            return func(*args, **kwargs)
+
+        return _wrapper
+
+    return _decorator

src/sparseml/deepsparse/framework/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+
+"""
+Functionality related to integrating with, detecting, and getting information for
+support and sparsification in the DeepSparse framework.
+"""
+
+# flake8: noqa
+
+from .info import *

src/sparseml/deepsparse/framework/info.py

@@ -0,0 +1,153 @@
+# 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.
+
+"""
+Functionality related to detecting and getting information for
+support and sparsification in the DeepSparse framework.
+"""
+
+import logging
+from typing import Any
+
+from sparseml.base import Framework, get_version
+from sparseml.deepsparse.base import check_deepsparse_install
+from sparseml.deepsparse.sparsification import sparsification_info
+from sparseml.framework import FrameworkInferenceProviderInfo, FrameworkInfo
+from sparseml.sparsification import SparsificationInfo
+from sparsezoo import File, Model
+
+
+__all__ = ["is_supported", "detect_framework", "framework_info"]
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def is_supported(item: Any) -> bool:
+    """
+    :param item: The item to detect the support for
+    :type item: Any
+    :return: True if the item is supported by deepsparse, False otherwise
+    :rtype: bool
+    """
+    framework = detect_framework(item)
+
+    return framework == Framework.deepsparse
+
+
+def detect_framework(item: Any) -> Framework:
+    """
+    Detect the supported ML framework for a given item specifically for the
+    deepsparse package.
+    Supported input types are the following:
+    - A Framework enum
+    - A string of any case representing the name of the framework
+      (deepsparse, onnx, keras, pytorch, tensorflow_v1)
+    - A supported file type within the framework such as model files:
+      (onnx, pth, h5, pb)
+    - An object from a supported ML framework such as a model instance
+    If the framework cannot be determined, will return Framework.unknown
+
+    :param item: The item to detect the ML framework for
+    :type item: Any
+    :return: The detected framework from the given item
+    :rtype: Framework
+    """
+    framework = Framework.unknown
+
+    if isinstance(item, Framework):
+        _LOGGER.debug("framework detected from Framework instance")
+        framework = item
+    elif isinstance(item, str) and item.lower().strip() in Framework.__members__:
+        _LOGGER.debug("framework detected from Framework string instance")
+        framework = Framework[item.lower().strip()]
+    elif isinstance(item, str) and (
+        "deepsparse" in item.lower().strip() or "deep sparse" in item.lower().strip()
+    ):
+        _LOGGER.debug("framework detected from deepsparse text")
+        # string, check if it's a string saying deepsparse first
+        framework = Framework.deepsparse
+    elif isinstance(item, str) and ".onnx" in item.lower().strip():
+        _LOGGER.debug("framework detected from .onnx")
+        # string, check if it's a file url or path that ends with onnx extension
+        framework = Framework.deepsparse
+    elif isinstance(item, Model) or isinstance(item, File):
+        _LOGGER.debug("framework detected from SparseZoo instance")
+        # sparsezoo model/file, deepsparse supports these natively
+        framework = Framework.deepsparse
+
+    return framework
+
+
+def framework_info() -> FrameworkInfo:
+    """
+    Detect the information for the deepsparse framework such as package versions,
+    availability for core actions such as training and inference,
+    sparsification support, and inference provider support.
+
+    :return: The framework info for deepsparse
+    :rtype: FrameworkInfo
+    """
+    arch = {}
+
+    if check_deepsparse_install(raise_on_error=False):
+        from deepsparse.cpu import cpu_architecture
+
+        arch = cpu_architecture()
+
+    cpu_warnings = []
+    if arch and arch.isa != "avx512":
+        cpu_warnings.append(
+            "AVX512 instruction set not detected, inference performance will be limited"
+        )
+    if arch and arch.isa != "avx512" and arch.isa != "avx2":
+        cpu_warnings.append(
+            "AVX2 and AVX512 instruction sets not detected, "
+            "inference performance will be severely limited"
+        )
+    if arch and not arch.vni:
+        cpu_warnings.append(
+            "VNNI instruction set not detected, "
+            "quantized inference performance will be limited"
+        )
+
+    cpu_provider = FrameworkInferenceProviderInfo(
+        name="cpu",
+        description=(
+            "Performant CPU provider within DeepSparse specializing in speedup of "
+            "sparsified models using AVX and VNNI instruction sets"
+        ),
+        device="cpu",
+        supported_sparsification=SparsificationInfo(),  # TODO: fill in when available
+        available=check_deepsparse_install(raise_on_error=False),
+        properties={
+            "cpu_architecture": arch,
+        },
+        warnings=cpu_warnings,
+    )
+
+    return FrameworkInfo(
+        framework=Framework.deepsparse,
+        package_versions={
+            "deepsparse": get_version(package_name="deepsparse", raise_on_error=False),
+            "sparsezoo": get_version(package_name="sparsezoo", raise_on_error=False),
+            "sparseml": get_version(package_name="sparseml", raise_on_error=False),
+        },
+        sparsification=sparsification_info(),
+        inference_providers=[cpu_provider],
+        training_available=False,
+        sparsification_available=False,
+        exporting_onnx_available=False,
+        inference_available=True,
+    )

src/sparseml/deepsparse/sparsification/__init__.py

@@ -0,0 +1,24 @@
+# flake8: noqa
+
+# 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.
+
+"""
+Functionality related to applying, describing, and supporting sparsification
+algorithms to models within in the DeepSparse framework.
+"""
+
+# flake8: noqa
+
+from .info import *

src/sparseml/deepsparse/sparsification/info.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""
+Functionality related to describing availability and information of sparsification
+algorithms to models within in the DeepSparse framework.
+"""
+
+import logging
+
+from sparseml.sparsification import SparsificationInfo
+
+
+__all__ = ["sparsification_info"]
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def sparsification_info() -> SparsificationInfo:
+    """
+    Load the available setup for sparsifying model within deepsparse.
+    :return: The sparsification info for the deepsparse framework
+    :rtype: SparsificationInfo
+    """
+    _LOGGER.debug("getting sparsification info for deepsparse")
+    info = SparsificationInfo(modifiers=[])
+    _LOGGER.info("retrieved sparsification info for deepsparse: %s", info)
+
+    return info