incorporate sparsity into size calculation

Author: Sara Adkins

src/sparsezoo/analyze_v2/memory_access_analysis.py

@@ -165,7 +165,8 @@ def get_memory_access_bits(
         node_weight = get_node_weight(model_graph, node)
         precision = get_numpy_quantization_level(node_weight)
         counts = memory_access_counts["single"]["counts"]
-        bits = counts * precision
+        counts_sparse = memory_access_counts["single"]["counts_sparse"]
+        bits = (counts - counts_sparse) * precision
         is_quantized = is_quantized_layer(model_graph, node)
 
         return {

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,

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,
             },
         }