diff --git a/setup.py b/setup.py index 10fbef99..e5bc4fe4 100644 --- a/setup.py +++ b/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", ] diff --git a/src/sparsezoo/__init__.py b/src/sparsezoo/__init__.py index 94eb243c..e9c69869 100644 --- a/src/sparsezoo/__init__.py +++ b/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") diff --git a/src/sparsezoo/analyze_cli.py b/src/sparsezoo/analyze_cli.py index 8afe0cc1..abcee04b 100644 --- a/src/sparsezoo/analyze_cli.py +++ b/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() diff --git a/src/sparsezoo/analyze_v2/model_analysis.py b/src/sparsezoo/analyze_v2/model_analysis.py index 85845b9c..328a1a01 100644 --- a/src/sparsezoo/analyze_v2/model_analysis.py +++ b/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: diff --git a/src/sparsezoo/utils/onnx/analysis.py b/src/sparsezoo/utils/onnx/analysis.py index f6e590a7..ace0e2cb 100644 --- a/src/sparsezoo/utils/onnx/analysis.py +++ b/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