Include Sparsity in Size Calculation (#463)

* initial fix * style * incorporate sparsity into size calculation * quality
neuralmagic · Feb 22, 2024 · 08f94c4 · 08f94c4
1 parent 36fd754
commit 08f94c4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/src/sparsezoo/analyze_v2/operation_analysis.py b/src/sparsezoo/analyze_v2/operation_analysis.py
@@ -167,9 +167,11 @@ def get_operation_bits(
             is_quantized_op = "32" not in str(precision)
 
             single_counts = ops["single"]["counts"]
-            single_bits = single_counts * precision
+            single_counts_sparse = ops["single"]["counts_sparse"]
+            single_bits = (single_counts - single_counts_sparse) * precision
             block4_counts = ops["block4"]["counts"]
-            block4_bits = block4_counts * precision
+            block4_counts_sparse = ops["block4"]["counts_sparse"]
+            block4_bits = (block4_counts - block4_counts_sparse) * precision
             return {
                 "tensor": {
                     "counts": single_counts,

diff --git a/src/sparsezoo/analyze_v2/parameter_analysis.py b/src/sparsezoo/analyze_v2/parameter_analysis.py
@@ -153,16 +153,17 @@ def get_parameter_bits(
     If the layer is quantized, assume all its elements in the ndarray
      are quantized
     """
-    num_weights, _, _ = get_node_param_counts(node, model_graph)
+    num_weights, num_bias, num_sparse_weights = get_node_param_counts(node, model_graph)
     if num_weights > 0:
         precision = get_node_weight_precision(model_graph, node)
         is_quantized = is_quantized_layer(model_graph, node)
+        num_non_sparse_weights = num_weights - num_sparse_weights + num_bias
         return {
             "tensor": {
                 "counts": num_weights,
                 "counts_quant": num_weights * is_quantized,
-                "bits": num_weights * precision,
-                "bits_quant": num_weights * precision * is_quantized,
+                "bits": num_non_sparse_weights * precision,
+                "bits_quant": num_non_sparse_weights * precision * is_quantized,
             },
         }