Add space complexity metrics

.pre-commit-config.yaml

@@ -22,7 +22,6 @@ repos:
     hooks:
       - id: isort
         name: isort (python)
-        args: ["--profile", "black"]
 
 
   - repo: https://github.com/myint/docformatter
@@ -35,4 +34,3 @@ repos:
     rev: 25.1.0
     hooks:
       - id: black
-        args: ["--line-length", "79"]

pyproject.toml

@@ -0,0 +1,7 @@
+# pyproject.toml
+[tool.black]
+line-length = 79
+
+[tool.isort]
+profile = "black"
+line_length = 79

src/configs/generate_config.py

@@ -79,7 +79,6 @@ def get_activations_and_quantizers(
             activations_and_quantizers.append(
                 (get_nested_attribute(layer, activation_attribute), quantizer)
             )
-        print(f"AQ: {activations_and_quantizers}")
         return activations_and_quantizers
 
     def set_quantize_activations(
@@ -88,7 +87,6 @@ def set_quantize_activations(
         for attribute, quantized_activation in zip(
             self.activations.keys(), quantize_activations
         ):
-            print(f"SA: {attribute} {quantized_activation}")
             set_nested_attribute(layer, attribute, quantized_activation)
 
     def get_output_quantizers(self, layer):

src/examples/data_analysis/generate_plots.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python3
+#!/usr/bin/env python3
 
 import argparse
 

src/examples/models/mlp.py

@@ -17,4 +17,14 @@
     }
 }
 
-qconfigs = {"qconfig": simple_qconfig}
+uniform_qconfig = {
+    "hidden": {
+        "weights": {"kernel": UniformQuantizer(bits=4, signed=True)},
+        "activations": {"activation": UniformQuantizer(bits=4, signed=False)},
+    }
+}
+
+qconfigs = {
+    "simple": simple_qconfig,
+    "uniform": uniform_qconfig,
+}

src/examples/run.py

@@ -9,6 +9,7 @@
 from tensorflow.keras.optimizers import Adam
 
 from configs.qmodel import apply_quantization
+from utils.metrics import compute_space_complexity_model
 
 
 def main(args):
@@ -49,9 +50,6 @@ def main(args):
         loss="categorical_crossentropy",
         metrics=["accuracy"],
     )
-    print(qmodel.summary())
-    print(f"qweights: {[w.name for w in qmodel.layers[1].weights]}")
-    # print(f"qactivations: {[w.name for w in qmodel.layers[1].weights]}")
 
     callback_tuples = [
         (CaptureWeightCallback(qlayer), qconfig[layer.name])
@@ -69,6 +67,14 @@ def main(args):
         callbacks=[callback for callback, _ in callback_tuples],
     )
 
+    qmodel(next(iter(test_dataset))[0])
+    space_complexity = compute_space_complexity_model(qmodel)
+    print(f"Space complexity: {space_complexity / 8 * 1/1024} kB")
+    original_space_complexity = compute_space_complexity_model(model)
+    print(
+        f"Original space complexity: {original_space_complexity / 8 * 1/1024} kB"
+    )
+
     output_dict = {}
     output_dict["global"] = hist.history
     for callback, qconfig in callback_tuples:

src/quantizers/flex_quantizer.py

@@ -31,6 +31,7 @@ def __init__(
         bits: int,
         n_levels: int,
         signed: bool = True,
+        name_suffix: str = "",
     ):
         """Constructor.
 
@@ -55,10 +56,12 @@ def __init__(
         self.levels = None  # possible output values
         self.thresholds = None  # boundaries between levels
 
+        self.name_suffix = name_suffix
+
     def build(self, tensor_shape, name: str, layer: tf.keras.layers.Layer):
 
         alpha = layer.add_weight(
-            "alpha",
+            name=f"{name}{self.name_suffix}_alpha",
             initializer=tf.keras.initializers.Constant(0.1),
             trainable=True,
             dtype=tf.float32,
@@ -68,7 +71,7 @@ def build(self, tensor_shape, name: str, layer: tf.keras.layers.Layer):
         self.alpha = alpha
 
         levels = layer.add_weight(
-            "levels",
+            name=f"{name}{self.name_suffix}_levels",
             initializer=tf.keras.initializers.Constant(
                 np.linspace(
                     min_value(self.alpha, self.signed),
@@ -84,7 +87,7 @@ def build(self, tensor_shape, name: str, layer: tf.keras.layers.Layer):
         self.levels = levels
 
         thresholds = layer.add_weight(
-            "thresholds",
+            name=f"{name}{self.name_suffix}_thresholds",
             initializer=tf.keras.initializers.Constant(
                 np.linspace(
                     min_value(self.alpha, self.signed),
@@ -112,13 +115,7 @@ def range(self):
     def delta(self):
         return self.range() / self.m_levels
 
-    @tf.custom_gradient
-    def quantize(self, x, alpha, levels, thresholds):
-        # Capture the values of the parameters
-        self.alpha = alpha
-        self.levels = levels
-        self.thresholds = thresholds
-
+    def quantize_op(self, x):
         # Quantize levels (uniform quantization)
         qlevels = self.delta() * tf.math.floor(self.levels / self.delta())
 
@@ -134,6 +131,19 @@ def quantize(self, x, alpha, levels, thresholds):
                 q,
             )
 
+        return q
+
+    @tf.custom_gradient
+    def quantize(self, x, alpha, levels, thresholds):
+        # Capture the values of the parameters
+        self.alpha = alpha
+        self.levels = levels
+        self.thresholds = thresholds
+
+        q = self.quantize_op(x)
+
+        qlevels = self.delta() * tf.math.floor(self.levels / self.delta())
+
         def grad(upstream):
             ##### dq_dx uses STE #####
             dq_dx = tf.where(

src/quantizers/uniform_quantizer.py

@@ -1,4 +1,4 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 """This module implements a uniform quantizer for quantizing weights and
 activations."""
@@ -12,6 +12,8 @@
     _QuantizeHelper,
 )
 
+from quantizers.common import delta, max_value, min_value, span
+
 
 class UniformQuantizer(_QuantizeHelper, Quantizer):
     """An uniform quantizer algorithm support both signed and unsigned
@@ -65,30 +67,45 @@ def __call__(self, w):
                 return tf.clip_by_value(w, tf.keras.backend.epsilon(), np.inf)
 
         alpha = layer.add_weight(
-            name.join("_alpha"),
+            name=f"{name}{self.name_suffix}_alpha",
             initializer=self.initializer,
             trainable=True,
             dtype=tf.float32,
             regularizer=self.regularizer,
             constraint=PositiveConstraint(),
         )
+        levels = layer.add_weight(
+            name=f"{name}{self.name_suffix}_levels",
+            trainable=False,
+            shape=(self.m_levels,),
+            dtype=tf.float32,
+        )
         self.alpha = alpha
-        return {"alpha": alpha}
+        self.levels = levels
+
+        return {"alpha": alpha, "levels": levels}
 
     def __call__(self, inputs, training, weights, **kwargs):
         return self.quantize(inputs, weights["alpha"])
 
     def range(self):
-        return 2 * self.alpha if self.signed else self.alpha
+        return span(self.alpha, self.signed)
 
     def delta(self):
-        return self.range() / self.m_levels
+        return delta(self.alpha, self.m_levels, self.signed)
 
-    def levels(self):
+    def compute_levels(self):
         """Compute the quantization levels."""
-        start = -self.alpha if self.signed else 0
+        start = min_value(self.alpha, self.signed)
         return tf.range(start, start + self.range(), self.delta())
 
+    def quantize_op(self, x):
+        clipped_x = tf.clip_by_value(x, self.levels[0], self.levels[-1])
+        delta_v = (
+            2 * self.alpha if self.signed else self.alpha
+        ) / self.m_levels
+        return delta_v * tf.math.floor(clipped_x / delta_v)
+
     @tf.custom_gradient
     def quantize(self, x, alpha):
         """Uniform quantization.
@@ -97,25 +114,22 @@ def quantize(self, x, alpha):
         :param alpha: alpha parameter
         :returns: quantized input tensor
         """
-        # Capture alpha
+        # Store alpha for other methods to use
         self.alpha = alpha
 
-        # Compute quantization levels
-        levels = self.levels()
-
-        # Clip input values between min and max levels (function is zero outside the range)
-        clipped_x = tf.clip_by_value(x, levels[0], levels[-1])
+        self.levels = self.compute_levels()
 
-        # Quantize input values
-        q = self.delta() * tf.math.floor(clipped_x / self.delta())
+        # Use direct parameter passing to avoid graph scope issues
+        q = self.quantize_op(x)
 
         def grad(upstream):
             # Gradient only flows through if the input is within range
-            ## Use STE to estimate the gradient
             dq_dx = tf.where(
                 tf.logical_and(
-                    tf.greater_equal(x, levels[0]),
-                    tf.less_equal(x, levels[-1]),
+                    tf.greater_equal(x, min_value(alpha, self.signed)),
+                    tf.less_equal(
+                        x, max_value(alpha, self.m_levels, self.signed)
+                    ),
                 ),
                 upstream,
                 tf.zeros_like(x),

src/quantizers/uniform_quantizer_test.py

@@ -52,7 +52,9 @@ def test_can_build_weights(self):
             name_suffix="_test",
         )
         weights = quantizer.build(self.input_shape, "test", self.mock_layer)
-        self.assertDictEqual(weights, {"alpha": weights["alpha"]})
+        self.assertDictEqual(
+            weights, {"alpha": weights["alpha"], "levels": weights["levels"]}
+        )
 
     # TODO(Fran): Consider using a fixture here?
     def assert_weights_within_limits(self, bits, signed):
@@ -71,7 +73,7 @@ def assert_weights_within_limits(self, bits, signed):
         output = quantizer(self.input_tensor, training=True, weights=weights)
 
         # Check that all output values are within the range determined by alpha
-        quantizer_levels = quantizer.levels()
+        quantizer_levels = quantizer.compute_levels()
         min = quantizer_levels[0]
         max = quantizer_levels[-1]
 
@@ -142,9 +144,9 @@ def test_expected_levels(self):
 
         quantizer.build(self.input_shape, "test", self.mock_layer)
 
-        levels = quantizer.levels()
+        levels = quantizer.compute_levels()
         expected_n_levels = 2**3
-        self.assertEqual(len(levels), expected_n_levels)
+        self.assertEqual(levels.shape.num_elements(), expected_n_levels)
 
         expected_levels = [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75]
         self.assertListEqual(list(levels), expected_levels)
@@ -161,7 +163,7 @@ def test_quantizer_levels_getitem(self):
 
         quantizer.build(self.input_shape, "test", self.mock_layer)
 
-        levels = quantizer.levels()
+        levels = quantizer.compute_levels()
         self.assertEqual(levels[0], -1.0)
         self.assertEqual(levels[2], -0.5)
         self.assertEqual(levels[7], 0.75)
@@ -192,7 +194,7 @@ def test_expected_levels_reflects_in_output_signed(self):
         # Call the quantizer
         output = quantizer(self.input_tensor, training=True, weights=weights)
         output_set = sorted(set(output.numpy().flatten()))
-        expected_set = list(quantizer.levels())
+        expected_set = list(quantizer.compute_levels())
 
         self.assertListEqual(output_set, expected_set)
 
@@ -221,7 +223,7 @@ def test_expected_levels_reflects_in_output_unsigned(self):
         # Call the quantizer
         output = quantizer(self.input_tensor, training=True, weights=weights)
         output_set = sorted(set(output.numpy().flatten()))
-        expected_set = list(quantizer.levels())
+        expected_set = list(quantizer.compute_levels())
 
         self.assertListEqual(output_set, expected_set)
 

src/utils/huffman.py

@@ -0,0 +1,39 @@
+import heapq
+from collections import Counter, namedtuple
+
+
+# A simple Node for the Huffman tree
+class Node(namedtuple("Node", ["freq", "value", "left", "right"])):
+    def __lt__(self, other):
+        return self.freq < other.freq
+
+
+def build_huffman_tree(weights):
+    counter = Counter(weights)
+    heap = [Node(freq, value, None, None) for value, freq in counter.items()]
+    heapq.heapify(heap)
+
+    while len(heap) > 1:
+        left = heapq.heappop(heap)
+        right = heapq.heappop(heap)
+        new_node = Node(left.freq + right.freq, None, left, right)
+        heapq.heappush(heap, new_node)
+
+    return heap[0]
+
+
+def assign_codes(node, prefix="", codebook={}):
+    if node.value is not None:
+        codebook[node.value] = prefix
+    else:
+        assign_codes(node.left, prefix + "0", codebook)
+        assign_codes(node.right, prefix + "1", codebook)
+    return codebook
+
+
+# TODO(frneer): Add tests for the Huffman tree
+# test_weights = [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
+# huffman_tree = build_huffman_tree(test_weights)
+# # print(huffman_tree)
+# huffman_codes = assign_codes(huffman_tree)
+# # print("Huffman Codes:", [len(code) for idx, code in huffman_codes.items()])