Print model analysis

src/sparsezoo/analyze/analysis.py

@@ -25,7 +25,6 @@
 from typing import Any, Dict, List, Optional, Union
 
 import numpy
-import onnx
 import yaml
 from onnx import ModelProto, NodeProto
 from pydantic import BaseModel, Field, PositiveFloat, PositiveInt
@@ -68,6 +67,7 @@
     is_parameterized_prunable_layer,
     is_quantized_layer,
     is_sparse_layer,
+    load_model,
 )
 
 
@@ -914,7 +914,7 @@ def from_onnx(cls, onnx_file_path: Union[str, ModelProto]):
             model_onnx = onnx_file_path
             model_name = ""
         else:
-            model_onnx = onnx.load(onnx_file_path)
+            model_onnx = load_model(onnx_file_path)
             model_name = str(onnx_file_path)
 
         model_graph = ONNXGraph(model_onnx)

src/sparsezoo/analyze_cli.py

@@ -77,6 +77,8 @@ def main(
         with open(save, "w") as file:
             file.write(analysis.to_yaml())
 
+    print(analysis)
+
 
 if __name__ == "__main__":
     main()

src/sparsezoo/analyze_v2/model_analysis.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from typing import Dict
+from typing import Dict, Optional
 
 import onnx
 import yaml
@@ -73,11 +73,70 @@ def to_dict(self):
             nodes=nodes,
         ).dict()
 
+    def calculate_sparsity_percentage(self, category: Dict):
+        counts_sparse = category["counts_sparse"]
+        counts = category["counts"]
+        return (counts_sparse / counts) * 100 if counts != 0 else 0
+
+    def calculate_quantized_percentage(self, tensor: Dict):
+        bits_quant = tensor["bits_quant"]
+        bits = tensor["bits"]
+        return (bits_quant / bits) * 100 if bits != 0 else 0
+
+    def __repr__(self):
+        data = self.to_dict()
+        summaries = data["summaries"]
+
+        param_total = summaries["params"]["sparsity"]["single"]["counts"]
+        param_sparsity = self.calculate_sparsity_percentage(
+            summaries["params"]["sparsity"]["single"]
+        )
+        param_size = summaries["params"]["quantization"]["tensor"]["bits"]
+        param_quantized = self.calculate_quantized_percentage(
+            summaries["params"]["quantization"]["tensor"]
+        )
+
+        ops_total = summaries["ops"]["sparsity"]["single"]["counts"]
+        ops_sparsity = self.calculate_sparsity_percentage(
+            summaries["ops"]["sparsity"]["single"]
+        )
+        ops_size = summaries["ops"]["quantization"]["tensor"]["bits"]
+        ops_quantized = self.calculate_quantized_percentage(
+            summaries["ops"]["quantization"]["tensor"]
+        )
+
+        mem_access_total = summaries["mem_access"]["sparsity"]["single"]["counts"]
+        mem_access_sparsity = self.calculate_sparsity_percentage(
+            summaries["mem_access"]["sparsity"]["single"]
+        )
+        mem_access_size = summaries["mem_access"]["quantization"]["tensor"]["bits"]
+        mem_access_quantized = self.calculate_quantized_percentage(
+            summaries["mem_access"]["quantization"]["tensor"]
+        )
+
+        return (
+            "Params:\n"
+            f"\ttotal\t\t: {param_total}\n"
+            f"\tsparsity%\t: {param_sparsity}\n"
+            f"\tsize [bits]\t: {param_size}\n"
+            f"\tquantized %\t: {param_quantized}\n"
+            "Ops:\n"
+            f"\ttotal\t\t: {ops_total}\n"
+            f"\tsparsity%\t: {ops_sparsity}\n"
+            f"\tsize [bits]\t: {ops_size}\n"
+            f"\tquantized %\t: {ops_quantized}\n"
+            "Memory Access:\n"
+            f"\ttotal\t\t: {mem_access_total}\n"
+            f"\tsparsity%\t: {mem_access_sparsity}\n"
+            f"\tsize [bits]\t: {mem_access_size}\n"
+            f"\tquantized %\t: {mem_access_quantized}\n"
+        )
+
     def to_yaml(self):
         return yaml.dump(self.to_dict())
 
 
-def analyze(path: str) -> "ModelAnalysis":
+def analyze(path: str, download_path: Optional[str] = None) -> "ModelAnalysis":
     """
     Entry point to run the model analysis.
 
@@ -89,7 +148,7 @@ def analyze(path: str) -> "ModelAnalysis":
     if path.endswith(".onnx"):
         onnx_model = load_model(path)
     elif is_stub(path):
-        model = Model(path)
+        model = Model(path, download_path)
         onnx_model_path = model.onnx_model.path
         onnx_model = onnx.load(onnx_model_path)
     else:

src/sparsezoo/model/model.py

@@ -15,6 +15,7 @@
 import logging
 import os
 import re
+from pathlib import Path
 from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import numpy
@@ -88,6 +89,7 @@ def __init__(self, source: str, download_path: Optional[str] = None):
             _setup_args = self.initialize_model_from_stub(stub=self.source)
             files, path, url, validation_results, compressed_size = _setup_args
             if download_path is not None:
+                download_path = str(Path(download_path).expanduser().resolve())
                 path = download_path  # overwrite cache path with user input
         else:
             # initializing the model from the path
@@ -703,23 +705,23 @@ def _download(
                 )
             ):
                 file.download(destination_path=directory_path)
-                return True
+                validations = True
             else:
                 _LOGGER.warning(
                     f"Failed to download file {file.name}. The url of the file is None."
                 )
-                return False
+                validations = False
 
         elif isinstance(file, Recipes):
-            validations = (
+            validations = all(
                 self._download(_file, directory_path) for _file in file.recipes
             )
 
         else:
-            validations = (
+            validations = all(
                 self._download(_file, directory_path) for _file in file.values()
             )
-            return all(validations)
+        return validations
 
     def _sample_outputs_list_to_dict(
         self,

src/sparsezoo/utils/onnx/external_data.py

@@ -174,7 +174,7 @@ def validate_onnx(model: Union[str, ModelProto]):
         raise ValueError(f"Invalid onnx model: {err}")
 
 
-def load_model(model: Union[str, ModelProto]) -> ModelProto:
+def load_model(model: Union[str, ModelProto, Path]) -> ModelProto:
     """
     Load an ONNX model from an onnx model file path. If a ModelProto
     is given, then it is returned.
@@ -185,7 +185,7 @@ def load_model(model: Union[str, ModelProto]) -> ModelProto:
     if isinstance(model, ModelProto):
         return model
 
-    if isinstance(model, str):
+    if isinstance(model, (Path, str)):
         return onnx.load(clean_path(model))
 
     raise ValueError(f"unknown type given for model: {type(model)}")