diff --git a/setup.cfg b/setup.cfg index 80621586b9..0dbc6fe79a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -5,7 +5,7 @@ ensure_newline_before_comments = True force_grid_wrap = 0 include_trailing_comma = True known_first_party = deepsparse,sparsezoo -known_third_party = numpy,onnx,requests,onnxruntime,flask,flask_cors +known_third_party = numpy,onnx,requests,onnxruntime,flask,flask_cors,tqdm sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER line_length = 88 diff --git a/setup.py b/setup.py index ea9a8594f9..2697b96e9f 100644 --- a/setup.py +++ b/setup.py @@ -38,7 +38,7 @@ binary_regexes = ["*/*.so", "*/*.so.*", "*.bin", "*/*.bin"] -_deps = ["numpy>=1.16.3", "onnx>=1.5.0,<1.8.0", "requests>=2.0.0"] +_deps = ["numpy>=1.16.3", "onnx>=1.5.0,<1.8.0", "requests>=2.0.0", "tqdm>=4.0.0"] _nm_deps = [f"{'sparsezoo-nightly' if _NIGHTLY else 'sparsezoo'}~={_VERSION}"] diff --git a/src/deepsparse/engine.py b/src/deepsparse/engine.py index a794b2ecd9..dbe909fdf0 100644 --- a/src/deepsparse/engine.py +++ b/src/deepsparse/engine.py @@ -23,13 +23,19 @@ import numpy from deepsparse.benchmark import BenchmarkResults +from tqdm.auto import tqdm try: + from sparsezoo import Zoo from sparsezoo.objects import File, Model -except Exception: + + sparsezoo_import_error = None +except Exception as sparsezoo_err: + Zoo = None Model = object File = object + sparsezoo_import_error = sparsezoo_err try: # flake8: noqa @@ -55,9 +61,13 @@ def _model_to_path(model: Union[str, Model, File]) -> str: if not model: raise ValueError("model must be a path, sparsezoo.Model, or sparsezoo.File") - if isinstance(model, str): - pass - elif Model is not object and isinstance(model, Model): + if isinstance(model, str) and model.startswith("zoo:"): + # load SparseZoo Model from stub + if sparsezoo_import_error is not None: + raise sparsezoo_import_error + model = Zoo.load_model_from_stub(model) + + if Model is not object and isinstance(model, Model): # default to the main onnx file for the model model = model.onnx_file.downloaded_path() elif File is not object and isinstance(model, File): @@ -105,8 +115,9 @@ class Engine(object): | # create an engine for batch size 1 on all available cores | engine = Engine("path/to/onnx", batch_size=1, num_cores=None) - :param model: Either a path to the model's onnx file, a sparsezoo Model object, - or a sparsezoo ONNX File object that defines the neural network + :param model: Either a path to the model's onnx file, a SparseZoo model stub + prefixed by 'zoo:', a SparseZoo Model object, or a SparseZoo ONNX File + object that defines the neural network :param batch_size: The batch size of the inputs to be used with the engine :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 @@ -324,6 +335,7 @@ def benchmark( num_warmup_iterations: int = 5, include_inputs: bool = False, include_outputs: bool = False, + show_progress: bool = False, ) -> BenchmarkResults: """ A convenience function for quickly benchmarking the instantiated model @@ -342,6 +354,7 @@ def benchmark( will be added to the results. Default is False :param include_outputs: If True, outputs from forward passes during benchmarking will be added to the results. Default is False + :param show_progress: If True, will display a progress bar. Default is False :return: the results of benchmarking """ # define data loader @@ -355,6 +368,7 @@ def _infinite_loader(): num_warmup_iterations=num_warmup_iterations, include_inputs=include_inputs, include_outputs=include_outputs, + show_progress=show_progress, ) def benchmark_loader( @@ -364,6 +378,7 @@ def benchmark_loader( num_warmup_iterations: int = 5, include_inputs: bool = False, include_outputs: bool = False, + show_progress: bool = False, ) -> BenchmarkResults: """ A convenience function for quickly benchmarking the instantiated model @@ -382,6 +397,7 @@ def benchmark_loader( will be added to the results. Default is False :param include_outputs: If True, outputs from forward passes during benchmarking will be added to the results. Default is False + :param show_progress: If True, will display a progress bar. Default is False :return: the results of benchmarking """ assert num_iterations >= 1 and num_warmup_iterations >= 0, ( @@ -391,13 +407,15 @@ def benchmark_loader( completed_iterations = 0 results = BenchmarkResults() + if show_progress: + progress_bar = tqdm(total=num_iterations) + while completed_iterations < num_warmup_iterations + num_iterations: for batch in loader: # run benchmark start = time.time() out = self.run(batch) end = time.time() - completed_iterations += 1 if completed_iterations >= num_warmup_iterations: # update results if warmup iterations are completed @@ -408,10 +426,17 @@ def benchmark_loader( inputs=batch if include_inputs else None, outputs=out if include_outputs else None, ) + if show_progress: + progress_bar.update(1) + + completed_iterations += 1 if completed_iterations >= num_warmup_iterations + num_iterations: break + if show_progress: + progress_bar.close() + return results def _validate_inputs(self, inp: List[numpy.ndarray]): @@ -453,8 +478,9 @@ def compile_model( Gives defaults of batch_size == 1 and num_cores == None (will use all physical cores available on a single socket). - :param model: Either a path to the model's onnx file, a sparsezoo Model object, - or a sparsezoo ONNX File object that defines the neural network + :param model: Either a path to the model's onnx file, a SparseZoo model stub + prefixed by 'zoo:', a SparseZoo Model object, or a SparseZoo ONNX File + object that defines the neural network :param batch_size: The batch size of the inputs to be used with the model :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 @@ -473,6 +499,7 @@ def benchmark_model( num_warmup_iterations: int = 5, include_inputs: bool = False, include_outputs: bool = False, + show_progress: bool = False, ) -> BenchmarkResults: """ Convenience function to benchmark a model in the DeepSparse Engine @@ -480,8 +507,9 @@ def benchmark_model( Gives defaults of batch_size == 1 and num_cores == None (will use all physical cores available on a single socket). - :param model: Either a path to the model's onnx file, a sparsezoo Model object, - or a sparsezoo ONNX File object that defines the neural network + :param model: Either a path to the model's onnx file, a SparseZoo model stub + prefixed by 'zoo:', a SparseZoo Model object, or a SparseZoo ONNX File + object that defines the neural network :param batch_size: The batch size of the inputs to be used with the model :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 @@ -498,12 +526,18 @@ def benchmark_model( will be added to the results. Default is False :param include_outputs: If True, outputs from forward passes during benchmarking will be added to the results. Default is False + :param show_progress: If True, will display a progress bar. Default is False :return: the results of benchmarking """ model = compile_model(model, batch_size, num_cores) return model.benchmark( - inp, num_iterations, num_warmup_iterations, include_inputs, include_outputs + inp, + num_iterations, + num_warmup_iterations, + include_inputs, + include_outputs, + show_progress, ) @@ -524,9 +558,9 @@ def analyze_model( Gives defaults of batch_size == 1 and num_cores == None (will use all physical cores available on a single socket). - :param model: Either a path to the model's onnx file, a sparsezoo Model object, - or a sparsezoo ONNX File object that defines the neural network - graph definition to analyze + :param model: Either a path to the model's onnx file, a SparseZoo model stub + prefixed by 'zoo:', a SparseZoo Model object, or a SparseZoo ONNX File + 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