[BugFix] Add analyze to init

setup.py

@@ -51,6 +51,7 @@
     "pandas>1.3",
     "py-machineid>=0.3.0",
     "geocoder>=1.38.0",
+    "onnxruntime>=1.0.0",
 ]
 _notebook_deps = ["ipywidgets>=7.0.0", "jupyter>=1.0.0"]
 
@@ -62,7 +63,6 @@
     "isort>=5.7.0",
     "pytest>=6.0.0",
     "wheel>=0.36.2",
-    "onnxruntime>=1.0.0",
     "matplotlib>=3.0.0",
 ]
 

src/sparsezoo/__init__.py

@@ -25,5 +25,6 @@
 from . import deployment_package as deployment_package_module
 from .deployment_package import *
 from .analytics import sparsezoo_analytics as _analytics
+from .analyze_v2 import analyze
 
 _analytics.send_event("python__init")

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/utils/onnx/analysis.py

@@ -385,12 +385,15 @@ def get_node_weight(
     :param node: node to which parameter belongs to
     :return: a numpy array of param value, None if not found
     """
-
     initializer_name = get_node_weight_name(model_graph, node)
     weight = get_initializer_value(model_graph, node, initializer_name)
+
     if initializer_name is not None and weight is None and node.op_type != "Gather":
         raise Exception(f"Parameter for {node.name} not found")
 
+    # some weights are not accessible, and returns the zero_points, which are scalars
+    if weight is not None and weight.ndim in [0, 1]:
+        return None
     return weight