[Export Refactor][Image Classification] apply_optimizations function

src/sparseml/core/session.py

@@ -100,7 +100,7 @@ def pre_initialize_structure(
         This will run the pre-initialize structure method for each modifier in the
         session's lifecycle. This will also set the session's state to the
         pre-initialized state. Takes care of cases when the model(s) structure
-        has been previosuly modified by a modifier.
+        has been previously modified by a modifier.
 
         :param model: the model to pre-initialize the structure for
         :param recipe: the recipe to use for the sparsification, can be a path to a

src/sparseml/export/__init__.py

@@ -0,0 +1,13 @@
+# 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.

src/sparseml/export.py → src/sparseml/export/export.py

@@ -14,56 +14,26 @@
 
 import logging
 from pathlib import Path
-from typing import Any, Callable, List, Optional, Union
-
-from pydantic import BaseModel, Field
-
+from typing import Any, List, Optional, Union
+
+from sparseml.export.helpers import apply_optimizations
+from sparseml.exporters import ExportTargets
+from sparseml.integration_helper_functions import (
+    IntegrationHelperFunctions,
+    infer_integration,
+)
 from sparseml.pytorch.opset import TORCH_DEFAULT_ONNX_OPSET
-from sparsezoo.utils.registry import RegistryMixin
 
 
 _LOGGER = logging.getLogger(__name__)
-AVAILABLE_DEPLOYMENT_TARGETS = ["deepsparse", "onnxruntime"]
-
-
-class IntegrationHelperFunctions(BaseModel, RegistryMixin):
-    """
-    Registry that maps integration names to helper functions
-    for creation/export/manipulation of models for a specific
-    integration.
-    """
-
-    create_model: Optional[Callable] = Field(
-        description="A function that creates a (sparse) "
-        "PyTorch model from a source path."
-    )
-    create_dummy_input: Optional[Callable] = Field(
-        description="A function that creates a dummy input "
-        "given a (sparse) PyTorch model."
-    )
-    export_model: Optional[Callable] = Field(
-        description="A function that exports a (sparse) PyTorch "
-        "model to an ONNX format appropriate for a "
-        "deployment target."
-    )
-    apply_optimizations: Optional[Callable] = Field(
-        description="A function that takes a set of "
-        "optimizations and applies them to an ONNX model."
-    )
-    export_sample_inputs_outputs: Optional[Callable] = Field(
-        description="A function that exports input/output samples given "
-        "a (sparse) PyTorch model."
-    )
-    create_deployment_folder: Optional[Callable] = Field(
-        description="A function that creates a "
-        "deployment folder for the exporter ONNX model"
-        "with the appropriate structure."
-    )
+AVAILABLE_DEPLOYMENT_TARGETS = [target.value for target in ExportTargets]
+ONNX_MODEL_NAME = "model.onnx"
 
 
 def export(
     source_path: Union[Path, str],
     target_path: Union[Path, str],
+    model_onnx_name: str = ONNX_MODEL_NAME,
     deployment_target: str = "deepsparse",
     integration: Optional[str] = None,
     sample_data: Optional[Any] = None,
@@ -80,17 +50,18 @@ def export(
     Export a PyTorch model to a deployment target specified by the `deployment_target`.
 
     The functionality follows a set of steps:
-    1. Create a PyTorch model from the source_path.
-    2. Create a dummy input for the model.
-    3. Export the model, using the precomputed dummy input, to an
-        ONNX format appropriate for the deployment target.
-    4. Apply optimizations to the exported model (optional).
+    1. Create a PyTorch model from the file located in source_path.
+    2. Create model dummy input.
+    3. Export the model to the format specified by the `deployment_target`.
+    4. (Optional) Apply optimizations to the exported model.
     5. Export sample inputs and outputs for the exported model (optional).
     6. Create a deployment folder for the exported model with the appropriate structure.
     7. Validate the correctness of the exported model (optional).
 
     :param source_path: The path to the PyTorch model to export.
     :param target_path: The path to save the exported model to.
+    :param model_onnx_name: The name of the exported model.
+        Defaults to ONNX_MODEL_NAME.
     :param deployment_target: The deployment target to export
         the model to. Defaults to 'deepsparse'.
     :param integration: The name of the integration to use for
@@ -132,17 +103,25 @@ def export(
         IntegrationHelperFunctions.load_from_registry(integration)
     )
 
-    model = helper_functions.create_model(source_path, device)
+    # for now, this code is not runnable, serves as a blueprint
+    model, auxiliary_items = helper_functions.create_model(
+        source_path, **kwargs  # noqa: F821
+    )
     sample_data = (
-        helper_functions.create_dummy_input(model, batch_size)
+        helper_functions.create_dummy_input(**auxiliary_items)
         if sample_data is None
         else sample_data
     )
     onnx_file_path = helper_functions.export_model(
-        model, sample_data, target_path, deployment_target, opset, single_graph_file
+        model, sample_data, target_path, deployment_target, opset
     )
 
-    helper_functions.apply_optimizations(onnx_file_path, graph_optimizations)
+    apply_optimizations(
+        onnx_file_path=onnx_file_path,
+        graph_optimizations=graph_optimizations,
+        available_graph_optimizations=helper_functions.graph_optimizations,
+        single_graph_file=single_graph_file,
+    )
 
     if export_sample_inputs_outputs:
         helper_functions.export_sample_inputs_outputs(model, target_path)
@@ -166,20 +145,6 @@ def export(
     )
 
 
-def infer_integration(source_path: Union[Path, str]) -> str:
-    """
-    Infer the integration to use for exporting the model from the source_path.
-    For example:
-        - for transformers model the integration
-            can be inferred from `config.json`
-        - for computer vision, the integration
-            can be inferred from the model architecture (`arch_key`)
-    :param source_path: The path to the PyTorch model to export.
-    :return: The name of the integration to use for exporting the model.
-    """
-    raise NotImplementedError
-
-
 def validate_correctness(deployment_path: Union[Path, str]):
     """
     Validate the correctness of the exported model.

src/sparseml/export/helpers.py

@@ -0,0 +1,107 @@
+# 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.
+
+from collections import OrderedDict
+from enum import Enum
+from pathlib import Path
+from typing import Callable, List, Union
+
+import onnx
+
+from sparsezoo.utils.onnx import save_onnx
+
+
+__all__ = ["apply_optimizations"]
+
+
+class GraphOptimizationOptions(Enum):
+    """
+    Holds the string names of the graph optimization options.
+    """
+
+    none = "none"
+    all = "all"
+
+
+def apply_optimizations(
+    onnx_file_path: Union[str, Path],
+    available_optimizations: OrderedDict[str, Callable],
+    target_optimizations: Union[str, List[str]] = GraphOptimizationOptions.all.value,
+    single_graph_file: bool = True,
+):
+    """
+    Apply optimizations to the graph of the ONNX model.
+
+    :param onnx_file_path: The path to the ONNX model file.
+    :param available_optimizations: The graph optimizations available
+        for the model. It is an ordered mapping from the string names
+        to functions that alter the model
+    :param target_optimizations: The name(s) of optimizations to apply.
+        It can be either a list of string name or a single string option
+        that specifies the set of optimizations to apply.
+        If is string, refer to the `GraphOptimizationOptions` enum
+        for the available options.
+    :param single_graph_file: Whether to save the optimized graph to a single
+        file or split it into multiple files. By default, it is True.
+    """
+    optimizations: List[Callable] = resolve_graph_optimizations(
+        available_optimizations=available_optimizations,
+        optimizations=target_optimizations,
+    )
+
+    onnx_model = onnx.load(onnx_file_path)
+
+    for optimization in optimizations:
+        onnx_model = optimization(onnx_model)
+
+    if single_graph_file:
+        save_onnx(onnx_model, onnx_file_path)
+        return
+
+    save_onnx_multiple_files(onnx_model)
+
+
+def resolve_graph_optimizations(
+    available_optimizations: OrderedDict[str, Callable],
+    optimizations: Union[str, List[str]],
+) -> List[Callable]:
+    """
+    Get the optimization functions to apply to the onnx model.
+
+    :param available_optimizations: The graph optimizations available
+        for the model. It is an ordered mapping from the string names
+        to functions that alter the model
+    :param optimizations: The name(s) of optimizations to apply.
+        It can be either a list of string name or a single string option
+        that specifies the set of optimizations to apply.
+        If is string, refer to the `GraphOptimizationOptions` enum
+        for the available options.
+    return The list of optimization functions to apply.
+    """
+    if isinstance(optimizations, str):
+        if optimizations == GraphOptimizationOptions.none.value:
+            return []
+        elif optimizations == GraphOptimizationOptions.all.value:
+            return list(available_optimizations.values())
+        else:
+            raise KeyError(f"Unknown graph optimization option: {optimizations}")
+    elif isinstance(optimizations, list):
+        return [available_optimizations[optimization] for optimization in optimizations]
+    else:
+        raise KeyError(f"Unknown graph optimization option: {optimizations}")
+
+
+# TODO: To discuss with @bfineran
+def save_onnx_multiple_files(*args, **kwargs):
+    raise NotImplementedError

src/sparseml/exporters/__init__.py

@@ -11,3 +11,14 @@
 # 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
+
+
+class ExportTargets(Enum):
+    """
+    Holds the names of the supported export targets
+    """
+
+    deepsparse = "deepsparse"
+    onnx = "onnx"

src/sparseml/integration_helper_functions.py

@@ -0,0 +1,109 @@
+# 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.
+
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+from pydantic import BaseModel, Field
+
+from sparsezoo.utils.registry import RegistryMixin
+
+
+__all__ = ["IntegrationHelperFunctions", "infer_integration"]
+
+
+class Integrations(Enum):
+    """
+    Holds the names of the available integrations.
+    """
+
+    image_classification = "image-classification"
+
+
+class IntegrationHelperFunctions(RegistryMixin, BaseModel):
+    """
+    Registry that maps names to helper functions
+    for creation/export/manipulation of models for a specific
+    integration.
+    """
+
+    create_model: Optional[
+        Callable[
+            Tuple[Union[str, Path], Optional[Dict[str, Any]]],
+            Tuple["torch.nn.Module", Dict[str, Any]],  # noqa F821
+        ]
+    ] = Field(
+        description="A function that takes: "
+        "- a source path to a PyTorch model "
+        "- (optionally) a dictionary of additional arguments"
+        "and returns: "
+        "- a (sparse) PyTorch model "
+        "- (optionally) a dictionary of additional arguments"
+    )
+    create_dummy_input: Optional[
+        Callable[..., "torch.Tensor"]  # noqa F821
+    ] = Field(  # noqa: F82
+        description="A function that takes: "
+        "- a dictionary of arguments"
+        "and returns: "
+        "- a dummy input for the model (a torch.Tensor) "
+    )
+    export: Optional[Callable[..., str]] = Field(
+        description="A function that takes: "
+        " - a (sparse) PyTorch model "
+        " - sample input data "
+        " - the path to save the exported model to "
+        " - the name to save the exported ONNX model as "
+        " - the deployment target to export to "
+        " - the opset to use for the export "
+        " - (optionally) a dictionary of additional arguments"
+        "and returns nothing"
+    )
+    graph_optimizations: Optional[Dict[str, Callable]] = Field(
+        description="A mapping from names to graph optimization functions "
+    )
+    export_sample_inputs_outputs: Optional[Callable] = Field(
+        description="A function that exports input/output samples given "
+        "a (sparse) PyTorch model."
+    )
+    create_deployment_folder: Optional[Callable] = Field(
+        description="A function that creates a "
+        "deployment folder for the exporter ONNX model"
+        "with the appropriate structure."
+    )
+
+
+def infer_integration(source_path: Union[Path, str]) -> str:
+    """
+    Infer the integration to use for exporting the model from the source_path.
+
+    :param source_path: The path to the PyTorch model to export.
+    :return: The name of the integration to use for exporting the model.
+    """
+    from sparseml.pytorch.image_classification.utils.helpers import (
+        is_image_classification_model,
+    )
+
+    if is_image_classification_model(source_path):
+        # import to register the image_classification integration helper functions
+        import sparseml.pytorch.image_classification.integration_helper_functions  # noqa F401
+
+        return Integrations.image_classification.value
+    else:
+        raise ValueError(
+            f"Could not infer integration from source_path: {source_path}."
+            f"Please specify an argument `integration` from one of"
+            f"the available integrations: {list(Integrations)}."
+        )