Generalize disabling batch size across engine interfaces

src/deepsparse/engine.py

@@ -17,12 +17,12 @@
 """
 
 import logging
+import os
 import time
 from enum import Enum
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy
-import onnx
 from tqdm.auto import tqdm
 
 from deepsparse.analytics import deepsparse_analytics as _analytics
@@ -105,9 +105,11 @@ def from_str(key: str):
             raise ValueError(f"unsupported Scheduler: {key}")
 
 
-def _validate_batch_size(batch_size: int) -> int:
-    if batch_size < 1:
-        raise ValueError("batch_size must be greater than 0")
+def _validate_batch_size(batch_size: Optional[int]) -> Optional[int]:
+    if batch_size is None or batch_size < 1:
+        _LOGGER.warn("batch_size < 1 so disabling batch size override")
+        os.environ["NM_DISABLE_BATCH_OVERRIDE"] = "1"
+        return None
 
     return batch_size
 
@@ -225,12 +227,11 @@ class BaseEngine(object):
     def construct(
         self,
         model: Union[str, "Model", "File"],
-        batch_size: int = 1,
-        num_cores: int = None,
-        num_streams: int = None,
-        scheduler: Scheduler = None,
-        input_shapes: List[List[int]] = None,
-        disable_batch_override: bool = False,
+        batch_size: Optional[int] = 1,
+        num_cores: Optional[int] = None,
+        num_streams: Optional[int] = None,
+        scheduler: Optional[Scheduler] = None,
+        input_shapes: Optional[List[List[int]]] = None,
         kv_cache_params: Optional[KVCacheParams] = None,
     ):
         _analytics.send_event("python__engine__init")
@@ -240,18 +241,16 @@ def construct(
         self._num_streams = _validate_num_streams(num_streams, self._num_cores)
         self._scheduler = _validate_scheduler(scheduler)
         self._input_shapes = input_shapes
-        self._disable_batch_override = disable_batch_override
         self._kv_cache_params = kv_cache_params
         self._cpu_avx_type = AVX_TYPE
         self._cpu_vnni = VNNI
 
     def construct_with_context(
         self,
         model: Union[str, "Model", "File"],
-        batch_size: int,
+        batch_size: Optional[int],
         context: Context,
-        input_shapes: List[List[int]] = None,
-        disable_batch_override: bool = False,
+        input_shapes: Optional[List[List[int]]] = None,
         kv_cache_params: Optional[KVCacheParams] = None,
     ):
         _analytics.send_event("python__engine__init")
@@ -261,7 +260,6 @@ def construct_with_context(
         self._num_streams = context.num_streams
         self._scheduler = _validate_scheduler(context.scheduler)
         self._input_shapes = input_shapes
-        self._disable_batch_override = disable_batch_override
         self._kv_cache_params = kv_cache_params
         self._cpu_avx_type = AVX_TYPE
         self._cpu_vnni = VNNI
@@ -295,23 +293,27 @@ class Engine(BaseEngine):
     def __init__(
         self,
         model: Union[str, "Model", "File"],
-        batch_size: int = 1,
+        batch_size: Optional[int] = 1,
         num_cores: int = None,
         num_streams: int = None,
         scheduler: Scheduler = None,
-        input_shapes: List[List[int]] = None,
+        input_shapes: Optional[List[List[int]]] = None,
     ):
         BaseEngine.construct(
             self, model, batch_size, num_cores, num_streams, scheduler, input_shapes
         )
 
+        # self._batch_size is allowed to be None to disable setting a batch size,
+        # but the engine needs to be passed an integer. The value is abitrary and ignored
+        engine_batch_size = self._batch_size if self._batch_size else 1
+
         if self._input_shapes:
             with override_onnx_input_shapes(
                 self._model_path, self._input_shapes
             ) as model_path:
                 self._eng_net = LIB.deepsparse_engine(
                     model_path,
-                    self._batch_size,
+                    engine_batch_size,
                     self._num_cores,
                     self._num_streams,
                     self._scheduler.value,
@@ -320,7 +322,7 @@ def __init__(
         else:
             self._eng_net = LIB.deepsparse_engine(
                 self._model_path,
-                self._batch_size,
+                engine_batch_size,
                 self._num_cores,
                 self._num_streams,
                 self._scheduler.value,
@@ -699,14 +701,13 @@ def _validate_inputs(self, inp: List[numpy.ndarray]):
             raise ValueError("inp must be a list, given {}".format(type(inp)))
 
         for arr in inp:
-            if not self._disable_batch_override:
-                if arr.shape[0] != self._batch_size:
-                    raise ValueError(
-                        (
-                            "array batch size of {} must match the batch size "
-                            "the model was instantiated with {}"
-                        ).format(arr.shape[0], self._batch_size)
-                    )
+            if self._batch_size and arr.shape[0] != self._batch_size:
+                raise ValueError(
+                    (
+                        "array batch size of {} must match the batch size "
+                        "the model was instantiated with {}"
+                    ).format(arr.shape[0], self._batch_size)
+                )
 
             if not arr.flags["C_CONTIGUOUS"]:
                 raise ValueError(
@@ -762,14 +763,13 @@ class DebugAnalysisEngine(Engine):
     def __init__(
         self,
         model: Union[str, "Model", "File"],
-        batch_size: int = 1,
-        num_cores: int = None,
-        scheduler: Scheduler = None,
+        batch_size: Optional[int] = 1,
+        num_cores: Optional[int] = None,
+        scheduler: Optional[Scheduler] = None,
         input_shapes: List[List[int]] = None,
         num_iterations: int = 20,
         num_warmup_iterations: int = 5,
         optimization_level: int = 1,
-        disable_batch_override: bool = False,
         imposed_as: Optional[float] = None,
         imposed_ks: Optional[float] = None,
         kv_cache_params: Optional[KVCacheParams] = None,
@@ -782,12 +782,15 @@ def __init__(
             None,
             scheduler,
             input_shapes,
-            disable_batch_override,
             kv_cache_params,
         )
 
         # Helper
         def make_engine(self, model_path):
+            # self._batch_size is allowed to be None to disable setting a batch size,
+            # but the engine needs to be passed an integer. The value is abitrary and ignored
+            engine_batch_size = self._batch_size if self._batch_size else 1
+
             if self._kv_cache_params:
                 self._kv_cache = LIB.kv_cache(
                     self._kv_cache_params.prev_num_tokens,
@@ -796,7 +799,7 @@ def make_engine(self, model_path):
 
                 self._eng_net = LIB.deepsparse_engine(
                     model_path,
-                    self._batch_size,
+                    engine_batch_size,
                     self._num_cores,
                     self._num_streams,
                     self._scheduler.value,
@@ -814,7 +817,7 @@ def make_engine(self, model_path):
 
                 self._eng_net = LIB.deepsparse_engine(
                     model_path,
-                    self._batch_size,
+                    engine_batch_size,
                     self._num_cores,
                     self._num_streams,
                     self._scheduler.value,
@@ -880,21 +883,25 @@ class MultiModelEngine(Engine):
     def __init__(
         self,
         model: Union[str, "Model", "File"],
-        batch_size: int,
+        batch_size: Optional[int],
         context: Context,
-        input_shapes: List[List[int]] = None,
+        input_shapes: Optional[List[List[int]]] = None,
     ):
         BaseEngine.construct_with_context(
             self, model, batch_size, context, input_shapes
         )
 
+        # self._batch_size is allowed to be None to disable setting a batch size,
+        # but the engine needs to be passed an integer. The value is abitrary and ignored
+        engine_batch_size = self._batch_size if self._batch_size else 1
+
         if self._input_shapes:
             with override_onnx_input_shapes(
                 self._model_path, self._input_shapes
             ) as model_path:
                 self._eng_net = LIB.deepsparse_engine(
                     model_path,
-                    self._batch_size,
+                    engine_batch_size,
                     self._num_cores,
                     self._num_streams,
                     self._scheduler.value,
@@ -903,7 +910,7 @@ def __init__(
         else:
             self._eng_net = LIB.deepsparse_engine(
                 self._model_path,
-                self._batch_size,
+                engine_batch_size,
                 self._num_cores,
                 self._num_streams,
                 self._scheduler.value,
@@ -913,11 +920,11 @@ def __init__(
 
 def compile_model(
     model: Union[str, "Model", "File"],
-    batch_size: int = 1,
-    num_cores: int = None,
-    num_streams: int = None,
-    scheduler: Scheduler = None,
-    input_shapes: List[List[int]] = None,
+    batch_size: Optional[int] = 1,
+    num_cores: Optional[int] = None,
+    num_streams: Optional[int] = None,
+    scheduler: Optional[Scheduler] = None,
+    input_shapes: Optional[List[List[int]]] = None,
 ) -> Engine:
     """
     Convenience function to compile a model in the DeepSparse Engine
@@ -952,16 +959,16 @@ def compile_model(
 def benchmark_model(
     model: Union[str, "Model", "File"],
     inp: List[numpy.ndarray],
-    batch_size: int = 1,
-    num_cores: int = None,
-    num_streams: int = None,
+    batch_size: Optional[int] = 1,
+    num_cores: Optional[int] = None,
+    num_streams: Optional[int] = None,
     num_iterations: int = 20,
     num_warmup_iterations: int = 5,
     include_inputs: bool = False,
     include_outputs: bool = False,
     show_progress: bool = False,
-    scheduler: Scheduler = None,
-    input_shapes: List[List[int]] = None,
+    scheduler: Optional[Scheduler] = None,
+    input_shapes: Optional[List[List[int]]] = None,
 ) -> BenchmarkResults:
     """
     Convenience function to benchmark a model in the DeepSparse Engine
@@ -1019,16 +1026,15 @@ def benchmark_model(
 def model_debug_analysis(
     model: Union[str, "Model", "File"],
     inp: List[numpy.ndarray],
-    batch_size: int = 1,
-    num_cores: int = None,
+    batch_size: Optional[int] = 1,
+    num_cores: Optional[int] = None,
     num_iterations: int = 20,
     num_warmup_iterations: int = 5,
     optimization_level: int = 1,
-    disable_batch_override: bool = False,
     imposed_as: Optional[float] = None,
     imposed_ks: Optional[float] = None,
-    scheduler: Scheduler = None,
-    input_shapes: List[List[int]] = None,
+    scheduler: Optional[Scheduler] = None,
+    input_shapes: Optional[List[List[int]]] = None,
     kv_cache_params: Optional[KVCacheParams] = None,
 ) -> dict:
     """
@@ -1044,7 +1050,8 @@ def model_debug_analysis(
         object that defines the neural network graph definition to analyze
     :param inp: The list of inputs to pass to the engine for analyzing inference.
         The expected order is the inputs order as defined in the ONNX graph.
-    :param batch_size: The batch size of the inputs to be used with the model
+    :param batch_size: The batch size of the inputs to be used with the model,
+        <1 disables it.
     :param num_cores: The number of physical cores to run the model on.
         Pass None or 0 to run on the max number of cores
         for the current machine; default None
@@ -1054,7 +1061,6 @@ def model_debug_analysis(
         before analyzing, default is 5
     :param optimization_level: The amount of graph optimizations to perform.
         The current choices are either 0 (minimal) or 1 (all), default is 1
-    :param disable_batch_override: Indicates whether disable_batch_override was used or not
     :param imposed_as: Imposed activation sparsity, defaults to None.
         Will force the activation sparsity from all ReLu layers in the graph
         to match this desired sparsity level (percentage of 0's in the tensor).
@@ -1077,7 +1083,6 @@ def model_debug_analysis(
         num_iterations=num_iterations,
         num_warmup_iterations=num_warmup_iterations,
         optimization_level=optimization_level,
-        disable_batch_override=disable_batch_override,
         imposed_as=imposed_as,
         imposed_ks=imposed_ks,
         kv_cache_params=kv_cache_params,

src/deepsparse/utils/onnx.py

@@ -25,6 +25,7 @@
 from onnx import ModelProto
 from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
 
+from deepsparse.utils import parse_input_shapes
 from deepsparse.utils.extractor import Extractor
 from sparsezoo.utils import save_onnx, validate_onnx
 
@@ -206,11 +207,16 @@ def generate_random_inputs(
         if batch_size is not None:
             in_shape[0] = batch_size
 
-        _LOGGER.info(
-            "Generating input '{}', type = {}, shape = {}".format(
-                external_input.name, numpy.dtype(elem_type).name, in_shape
-            )
+        input_string = "input '{}', type = {}, shape = {}".format(
+            external_input.name, numpy.dtype(elem_type).name, in_shape
+        )
+
+        assert not any(dim < 1 for dim in in_shape), (
+            f"Dynamic shape found in {input_string}. "
+            "All shapes must be non-zero in order to generate random data"
         )
+
+        _LOGGER.info(f"Generating {input_string}")
         input_data_list.append(numpy.random.rand(*in_shape).astype(elem_type))
     return input_data_list
 
@@ -262,7 +268,7 @@ def override_onnx_batch_size(
 @contextlib.contextmanager
 def override_onnx_input_shapes(
     onnx_filepath: str,
-    input_shapes: Union[List[int], List[List[int]]],
+    input_shapes: Union[None, str, List[int], List[List[int]]],
     inplace: bool = True,
 ) -> str:
     """
@@ -291,6 +297,9 @@ def override_onnx_input_shapes(
         input for input in all_inputs if input.name not in initializer_input_names
     ]
 
+    if isinstance(input_shapes, str):
+        input_shapes = parse_input_shapes(input_shapes)
+
     # Input shapes should be a list of lists, even if there is only one input
     if not all(isinstance(inp, list) for inp in input_shapes):
         input_shapes = [input_shapes]