[Pipeline Refactor] Migration

src/deepsparse/__init__.py

@@ -33,7 +33,6 @@
 from .engine import *
 from .tasks import *
 from .pipeline import *
-from .base_pipeline import *
 from .loggers import *
 from .version import __version__, is_release
 from .analytics import deepsparse_analytics as _analytics

src/deepsparse/clip/decoder_pipeline.py

@@ -17,7 +17,7 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
-from deepsparse import Pipeline
+from deepsparse.legacy import Pipeline
 from deepsparse.utils import model_to_path
 
 

src/deepsparse/clip/text_pipeline.py

@@ -17,7 +17,7 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.utils import model_to_path
 from open_clip.tokenizer import tokenize
 

src/deepsparse/clip/visual_pipeline.py

@@ -22,7 +22,7 @@
 from torchvision.transforms import InterpolationMode
 
 from deepsparse.clip.constants import CLIP_RGB_MEANS, CLIP_RGB_STDS
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.pipelines.computer_vision import ComputerVisionSchema
 from deepsparse.utils import model_to_path
 

src/deepsparse/clip/zeroshot_pipeline.py

@@ -19,7 +19,7 @@
 from pydantic import BaseModel, Field
 
 from deepsparse.clip import CLIPTextInput, CLIPVisualInput
-from deepsparse.pipeline import BasePipeline, Pipeline
+from deepsparse.legacy.pipeline import BasePipeline, Pipeline
 from scipy.special import softmax
 
 

src/deepsparse/evaluation/cli.py

@@ -77,7 +77,11 @@
 )
 from src.deepsparse.evaluation.results import Result, save_evaluations
 from src.deepsparse.evaluation.utils import args_to_dict, get_save_path
-from src.deepsparse.pipeline import DEEPSPARSE_ENGINE, ORT_ENGINE, TORCHSCRIPT_ENGINE
+from src.deepsparse.operators.engine_operator import (
+    DEEPSPARSE_ENGINE,
+    ORT_ENGINE,
+    TORCHSCRIPT_ENGINE,
+)
 
 
 _LOGGER = logging.getLogger(__name__)

src/deepsparse/evaluation/evaluator.py

@@ -19,7 +19,11 @@
 )
 from src.deepsparse.evaluation.registry import EvaluationRegistry
 from src.deepsparse.evaluation.results import Result
-from src.deepsparse.pipeline import DEEPSPARSE_ENGINE, ORT_ENGINE, TORCHSCRIPT_ENGINE
+from src.deepsparse.operators.engine_operator import (
+    DEEPSPARSE_ENGINE,
+    ORT_ENGINE,
+    TORCHSCRIPT_ENGINE,
+)
 
 
 __all__ = ["evaluate"]

src/deepsparse/evaluation/utils.py

@@ -17,7 +17,8 @@
 
 from transformers import AutoModelForCausalLM
 
-from deepsparse import DEEPSPARSE_ENGINE, ORT_ENGINE, Pipeline
+from deepsparse import Pipeline
+from deepsparse.operators.engine_operator import DEEPSPARSE_ENGINE, ORT_ENGINE
 
 
 __all__ = ["text_generation_model_from_target", "get_save_path", "args_to_dict"]

src/deepsparse/image_classification/__init__.py

@@ -25,9 +25,12 @@
         "Please install deepsparse[image_classification] to use this pathway"
     )
 
-
 from .constants import *
-from .pipelines import *
+from .pipeline import *
+
+# flake8: noqa
+from .postprocess_operator import *
+from .preprocess_operator import *
 from .schemas import *
 
 

src/deepsparse/v2/image_classification/pipeline.py → src/deepsparse/image_classification/pipeline.py

@@ -13,40 +13,43 @@
 # limitations under the License.
 
 import logging
-import warnings
 from typing import Dict, Optional, Tuple, Union
 
-from deepsparse.v2.image_classification.postprocess_operator import (
+from deepsparse.image_classification.postprocess_operator import (
     ImageClassificationPostProcess,
 )
-from deepsparse.v2.image_classification.preprocess_operator import (
+from deepsparse.image_classification.preprocess_operator import (
     ImageClassificationPreProcess,
 )
-from deepsparse.v2.operators.engine_operator import EngineOperator
-from deepsparse.v2.pipeline import Pipeline
-from deepsparse.v2.routers.router import LinearRouter
-from deepsparse.v2.schedulers.scheduler import OperatorScheduler
+from deepsparse.operators.engine_operator import EngineOperator
+from deepsparse.operators.registry import OperatorRegistry
+from deepsparse.pipeline import Pipeline
+from deepsparse.routers.router import LinearRouter
+from deepsparse.schedulers.scheduler import OperatorScheduler
 
 
 _LOGGER = logging.getLogger(__name__)
 
 __all__ = ["ImageClassificationPipeline"]
 
 
+@OperatorRegistry.register(name="image_classification")
 class ImageClassificationPipeline(Pipeline):
     def __init__(
         self,
         model_path: str,
-        engine_kwargs: Optional[Dict] = None,
         class_names: Union[None, str, Dict[str, str]] = None,
         image_size: Optional[Tuple[int]] = None,
         top_k: int = 1,
+        **engine_kwargs,
     ):
+
         if not engine_kwargs:
             engine_kwargs = {}
             engine_kwargs["model_path"] = model_path
         elif engine_kwargs.get("model_path") != model_path:
-            warnings.warn(f"Updating engine_kwargs to include {model_path}")
+            _LOGGER.warning(f"Updating engine_kwargs to include {model_path}")
+            engine_kwargs["model_path"] = model_path
 
         engine = EngineOperator(**engine_kwargs)
         preproces = ImageClassificationPreProcess(

src/deepsparse/v2/image_classification/postprocess_operator.py → src/deepsparse/image_classification/postprocess_operator.py

@@ -18,7 +18,7 @@
 import numpy
 from pydantic import BaseModel, Field
 
-from deepsparse.v2.operators import Operator
+from deepsparse.operators import Operator
 
 
 class ImageClassificationOutput(BaseModel):

src/deepsparse/v2/image_classification/preprocess_operator.py → src/deepsparse/image_classification/preprocess_operator.py

@@ -23,8 +23,8 @@
     IMAGENET_RGB_MEANS,
     IMAGENET_RGB_STDS,
 )
+from deepsparse.operators import Operator
 from deepsparse.pipelines.computer_vision import ComputerVisionSchema
-from deepsparse.v2.operators import Operator
 
 
 class ImageClassificationInput(ComputerVisionSchema):

src/deepsparse/image_classification/validation_script.py

@@ -27,10 +27,6 @@
                                   on Imagenette  [default: zoo:cv/classificati
                                   on/resnet_v1-50/pytorch/sparseml/imagenette/
                                   base-none]
-  --batch-size, --batch_size INTEGER
-                                  Test batch size, must divide the dataset
-                                  evenly, else last batch will be dropped
-                                  [default: 1]
   --image-size, --image_size INTEGER
                                   integer size to evaluate images at (will be
                                   reshaped to square shape)  [default: 224]
@@ -111,15 +107,6 @@ def parse_json_callback(ctx, params, value: str) -> Dict:
     "evaluated. Defaults to dense (vanilla) resnet50 trained on Imagenette",
     show_default=True,
 )
-@click.option(
-    "--batch-size",
-    "--batch_size",
-    type=int,
-    default=1,
-    show_default=True,
-    help="Test batch size, must divide the dataset evenly, else last "
-    "batch will be dropped",
-)
 @click.option(
     "--image-size",
     "--image_size",
@@ -155,7 +142,6 @@ def parse_json_callback(ctx, params, value: str) -> Dict:
 def main(
     dataset_path: str,
     model_path: str,
-    batch_size: int,
     image_size: int,
     num_cores: int,
     dataset_kwargs: Dict,
@@ -206,18 +192,16 @@ def main(
 
     data_loader = DataLoader(
         dataset=dataset,
-        batch_size=batch_size,
+        batch_size=1,
         drop_last=True,
     )
 
     pipeline = Pipeline.create(
         task="image_classification",
         model_path=model_path,
-        batch_size=batch_size,
-        num_cores=num_cores,
         engine_type=engine,
+        num_cores=num_cores,
     )
-    print(f"engine info: {pipeline.engine}")
     correct = total = 0
     progress_bar = tqdm(data_loader)
 

src/deepsparse/v2/utils/__init__.py → src/deepsparse/legacy/__init__.py

@@ -1,5 +1,3 @@
-# flake8: noqa
-
 # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,9 +11,9 @@
 # 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 .helpers import *
-from .state import *
-from .types import *
 
+# flake8: noqa
 
-from .data import *  # isort:skip
+from .base_pipeline import *
+from .pipeline import *
+from .tasks import *

src/deepsparse/base_pipeline.py → src/deepsparse/legacy/base_pipeline.py

@@ -19,10 +19,10 @@
 from pydantic import BaseModel
 
 from deepsparse import Context
+from deepsparse.legacy.tasks import SupportedTasks, dynamic_import_task
 from deepsparse.loggers.base_logger import BaseLogger
 from deepsparse.loggers.build_logger import logger_from_config
 from deepsparse.loggers.constants import validate_identifier
-from deepsparse.tasks import SupportedTasks, dynamic_import_task
 
 
 __all__ = [
@@ -166,7 +166,7 @@ def create(
             implementation
         :return: pipeline object initialized for the given task
         """
-        from deepsparse.pipeline import Bucketable, BucketingPipeline, Pipeline
+        from deepsparse.legacy.pipeline import Bucketable, BucketingPipeline, Pipeline
 
         pipeline_constructor = BasePipeline._get_task_constructor(task)
         model_path = kwargs.get("model_path", None)
@@ -278,7 +278,7 @@ def from_config(
             logging. Default is None
         :return: loaded Pipeline object from the config
         """
-        from deepsparse.pipeline import PipelineConfig
+        from deepsparse.legacy.pipeline import PipelineConfig
 
         if isinstance(config, Path) or (
             isinstance(config, str) and os.path.exists(config)
@@ -308,7 +308,7 @@ def to_config(self) -> "PipelineConfig":  # noqa: F821
         """
         :return: PipelineConfig that can be used to reload this object
         """
-        from deepsparse.pipeline import PipelineConfig
+        from deepsparse.legacy.pipeline import PipelineConfig
 
         if not hasattr(self, "task"):
             raise RuntimeError(

src/deepsparse/image_classification/pipelines.py → src/deepsparse/legacy/image_classification/pipelines.py

@@ -31,7 +31,7 @@
     ImageClassificationInput,
     ImageClassificationOutput,
 )
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.utils import model_to_path