add tests for combining weights from activation and weight nodes in virtual node construction

irenab · irenab · commit 8ec2fe6a212c · 2025-03-24T15:14:20.000+02:00
diff --git a/model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py b/model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py
@@ -140,20 +140,21 @@ def __init__(self,
         """
         # Validate weights node
         kernel_attrs = fw_info.get_kernel_op_attributes(weights_node.type)
-        assert len(kernel_attrs) == 1 and kernel_attrs[0] is not None, 'Expected exactly one kernel attr.'
+        assert len(kernel_attrs) == 1 and kernel_attrs[0] is not None, f'Expected exactly one kernel attr, {kernel_attrs}'
         kernel_attr = kernel_attrs[0]
         conf_weights = [attr for attr in weights_node.weights if weights_node.is_configurable_weight(attr)]
         if len(conf_weights) > 1 or len(conf_weights) == 1 and not weights_node.is_configurable_weight(kernel_attr):
-            raise NotImplementedError('Only kernel weight can be configurable.')    # pragma: no cover
+            raise NotImplementedError(f'Only kernel weight can be configurable. Got configurable {conf_weights}.')
 
-        weights = weights_node.weights
+        weights = weights_node.weights.copy()
         act_node_w_rename = {}
         if act_node.weights:
-            assert fw_info.get_kernel_op_attributes(act_node)[0] is None, \
-                f'Node {act_node} with kernel cannot be used as activation for VirtualActivationWeightsNode.'
+            if not fw_info.get_kernel_op_attributes(act_node)[0] is None:
+                raise NotImplementedError(f'Node {act_node} with kernel cannot be used as activation for '
+                                          f'VirtualActivationWeightsNode.')
             if act_node.has_any_configurable_weight():
-                raise NotImplementedError('Node with a configurable weight cannot be used as activation for '
-                                          'VirtualActivationWeightsNode.')    # pragma: no cover
+                raise NotImplementedError(f'Node {act_node} with a configurable weight cannot be used as activation for '
+                                          'VirtualActivationWeightsNode.')
             # combine weights from activation and weights
             for w_id, w in act_node.weights.items():
                 if w_id not in weights and not (isinstance(w_id, str) and kernel_attr in w_id):
diff --git a/tests_pytest/_test_util/graph_builder_utils.py b/tests_pytest/_test_util/graph_builder_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Union, Iterable
+from typing import Union, Iterable, List
 
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.core import QuantizationConfig
@@ -31,13 +31,20 @@ class DummyLayer:
     pass
 
 
-def build_node(name='node', canonical_weights: dict=None, qcs=None, input_shape=(4, 5, 6), output_shape=(4, 5, 6),
+def build_node(name='node', canonical_weights: dict = None, final_weights: dict = None,
+               qcs: List[CandidateNodeQuantizationConfig] = None,
+               input_shape=(4, 5, 6), output_shape=(4, 5, 6),
                layer_class=DummyLayer, reuse=False):
     """ Build a node for tests.
-        Canonical weights are converted into full unique names.
+        Either 'canonical_weights' (to be used by default) or 'final_weights' should be passed.
+          Canonical weights are converted into full unique names, that contain the canonical name as a substring.
+          Final weights are used as is.
         candidate_quantization_cfg is set is qcs is passed."""
-    weights = canonical_weights or {}
-    weights = {k if isinstance(k, int) else full_attr_name(k): w for k, w in weights.items()}
+    assert canonical_weights is None or final_weights is None
+    if canonical_weights:
+        weights = {k if isinstance(k, int) else full_attr_name(k): w for k, w in canonical_weights.items()}
+    else:
+        weights = final_weights or {}
     node = BaseNode(name=name,
                     framework_attr={},
                     input_shape=input_shape,
@@ -46,6 +53,7 @@ def build_node(name='node', canonical_weights: dict=None, qcs=None, input_shape=
                     layer_class=layer_class,
                     reuse=reuse)
     if qcs:
+        assert isinstance(qcs, list)
         node.candidates_quantization_cfg = qcs
     return node
 
@@ -61,10 +69,24 @@ def full_attr_name(canonical_name: Union[str, dict, Iterable]):
     return canonical_name.__class__([convert(name) for name in canonical_name])
 
 
-def build_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, ())):
-    """ Build quantization config for tests.
-        w_attr contains {canonical name: (nbits, q_enabled)}
-        pos_attr: (nbits, q enabled, indices) """
+def build_nbits_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, ()),
+                   convert_canonical_attr=True) -> CandidateNodeQuantizationConfig:
+    """
+    Build quantization config with configurable nbits and enabling/disabling quantization only.
+
+    Args:
+        a_nbits: activation num bits.
+        a_enable: whether to enable activation quantization.
+        w_attr: quantization configuration for weight attributes in format {canonical name: (nbits, q_enabled)}.
+          By default, a canonical weight name is expected and is automatically converted to a dummy full name (that
+          contains the canonical name as a substring).
+          Final name can be passed along with convert_canonical_attr=False.
+        pos_attr: quantization configuration for positional weights in format (nbits, q enabled, indices).
+        convert_canonical_attr: whether to convert w_attr keys to full names.
+
+    Returns:
+
+    """
     w_attr = w_attr or {}
     attr_weights_configs_mapping = {
         k: AttributeQuantizationConfig(weights_n_bits=v[0], enable_weights_quantization=v[1])
@@ -91,7 +113,9 @@ def build_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, ())):
                                               activation_quantization_fn=None,
                                               activation_quantization_params_fn=None)
     # full names from the layers
-    attr_names = [full_attr_name(k) for k in w_attr.keys()]
+    attr_names = list(w_attr.keys())
+    if convert_canonical_attr:
+        attr_names = [full_attr_name(k) for k in w_attr.keys()]
     w_qcfg = NodeWeightsQuantizationConfig(qc=qc, op_cfg=op_cfg,
                                            weights_channels_axis=None,
                                            node_attrs_list=attr_names + list(pos_attr[2]))
diff --git a/tests_pytest/common_tests/unit_tests/core/graph/__init__.py b/tests_pytest/common_tests/unit_tests/core/graph/__init__.py
@@ -0,0 +1,14 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
diff --git a/tests_pytest/common_tests/unit_tests/core/graph/test_virtual_activation_weights_node.py b/tests_pytest/common_tests/unit_tests/core/graph/test_virtual_activation_weights_node.py
@@ -0,0 +1,129 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import numpy as np
+import pytest
+
+from model_compression_toolkit.core.common.graph.virtual_activation_weights_node import VirtualActivationWeightsNode
+from tests_pytest._test_util.graph_builder_utils import build_node, DummyLayer, build_nbits_qc
+
+
+class DummyLayerWKernel:
+    pass
+
+
+class TestVirtualActivationWeightsNode:
+    # TODO tests only cover combining weights from activation and weight nodes and errors.
+    def test_activation_with_weights(self, fw_info_mock):
+        """ Tests that weights from activation and weight node are combined correctly. """
+        # Each node has a unique weight attr and a unique positional weights. In addition, both nodes have
+        # an identical canonical attribute (but different full name), and an identical positional weight.
+        # All weights have different quantization.
+        a_node = build_node('a', layer_class=DummyLayer,
+                            final_weights={'aaweightaa': np.ones((3, 14)), 'foo': np.ones(15),
+                                           1: np.ones(15), 2: np.ones((5, 9))},
+                            qcs=[build_nbits_qc(a_nbits=5,
+                                                w_attr={'aaweightaa': (2, True), 'foo': (3, True)},
+                                                pos_attr=(4, True, [1, 2]),
+                                                convert_canonical_attr=False)])
+        w_node = build_node('w', layer_class=DummyLayerWKernel,
+                            final_weights={'wwweightww': np.ones((2, 71)), 'bar': np.ones(8),
+                                           1: np.ones(28), 3: np.ones(18)},
+                            qcs=[build_nbits_qc(a_nbits=6,
+                                                w_attr={'wwweightww': (5, True), 'bar': (6, True)},
+                                                pos_attr=(7, True, [1, 3]),
+                                                convert_canonical_attr=False)])
+
+        fw_info_mock.get_kernel_op_attributes = lambda nt: ['weight'] if nt is DummyLayerWKernel else [None]
+
+        v_node = VirtualActivationWeightsNode(a_node, w_node, fw_info_mock)
+        assert len(v_node.weights) == 8
+
+        assert len(w_node.weights) == len(a_node.weights) == 4
+        # weights from weight node are unchanged
+        for k, v in w_node.weights.items():
+            assert np.array_equal(v_node.weights.pop(k), v)
+        # unique weights from activation node are unchanged
+        assert np.array_equal(v_node.weights.pop('foo'), a_node.weights['foo'])
+        assert np.array_equal(v_node.weights.pop(2), a_node.weights[2])
+        # duplicate positional weight
+        assert np.array_equal(v_node.weights.pop(101), a_node.weights[1])
+        # duplicate weight attribute
+        [(new_attr, w)] = v_node.weights.items()
+        assert 'weight' not in new_attr
+        assert np.array_equal(w, a_node.weights['aaweightaa'])
+
+        assert len(v_node.candidates_quantization_cfg) == 1
+        v_qc = v_node.candidates_quantization_cfg[0]
+        v_attr_cfg = v_qc.weights_quantization_cfg.attributes_config_mapping
+        v_pos_cfg = v_qc.weights_quantization_cfg.pos_attributes_config_mapping
+        a_qc = a_node.candidates_quantization_cfg[0]
+        w_qc = w_node.candidates_quantization_cfg[0]
+
+        assert v_attr_cfg == {
+            'wwweightww': w_qc.weights_quantization_cfg.attributes_config_mapping['wwweightww'],
+            'bar': w_qc.weights_quantization_cfg.attributes_config_mapping['bar'],
+            'foo': a_qc.weights_quantization_cfg.attributes_config_mapping['foo'],
+            new_attr: a_qc.weights_quantization_cfg.attributes_config_mapping['aaweightaa']
+        }
+        assert v_pos_cfg == {
+            1: w_qc.weights_quantization_cfg.pos_attributes_config_mapping[1],
+            101: a_qc.weights_quantization_cfg.pos_attributes_config_mapping[1],
+            2: a_qc.weights_quantization_cfg.pos_attributes_config_mapping[2],
+            3: w_qc.weights_quantization_cfg.pos_attributes_config_mapping[3]
+        }
+
+    def test_invalid_configurable_w_node_weight(self, fw_info_mock):
+        w_node = build_node('w', layer_class=DummyLayerWKernel,
+                            canonical_weights={'kernel': np.ones(3), 'foo': np.ones(14)},
+                            qcs=[
+                                build_nbits_qc(w_attr={'kernel': (8, True), 'foo': (8, True)}),
+                                build_nbits_qc(w_attr={'kernel': (8, True), 'foo': (4, True)})
+                            ])
+        a_node = build_node('a', qcs=[build_nbits_qc()])
+
+        fw_info_mock.get_kernel_op_attributes = lambda nt: ['kernel'] if nt is DummyLayerWKernel else [None]
+
+        with pytest.raises(NotImplementedError, match='Only kernel weight can be configurable. Got configurable .*foo'):
+            VirtualActivationWeightsNode(a_node, w_node, fw_info_mock)
+
+    def test_invalid_a_node_configurable_weight(self, fw_info_mock):
+        w_node = build_node('w', layer_class=DummyLayerWKernel,
+                            canonical_weights={'kernel': np.ones(3), 'foo': np.ones(14)},
+                            qcs=[
+                                build_nbits_qc(w_attr={'kernel': (8, True), 'foo': (8, True)}),
+                                build_nbits_qc(w_attr={'kernel': (4, True), 'foo': (8, True)})
+                            ])
+        a_node = build_node('aaa', canonical_weights={'bar': np.ones(3), 'baz': np.ones(14)},
+                            qcs=[
+                                build_nbits_qc(w_attr={'bar': (8, True), 'baz': (8, True)}),
+                                build_nbits_qc(w_attr={'bar': (8, True), 'baz': (4, True)})
+                            ])
+        fw_info_mock.get_kernel_op_attributes = lambda nt: ['kernel'] if nt is DummyLayerWKernel else [None]
+
+        with pytest.raises(NotImplementedError, match='Node .*aaa with a configurable weight cannot be used as '
+                                                      'activation for VirtualActivationWeightsNode'):
+            VirtualActivationWeightsNode(a_node, w_node, fw_info_mock)
+
+    def test_invalid_a_node_kernel(self, fw_info_mock):
+        w_node = build_node('w', layer_class=DummyLayerWKernel, canonical_weights={'weight': np.ones(3)},
+                            qcs=[build_nbits_qc(w_attr={'weight': (8, True)})])
+        a_node = build_node('aaa', canonical_weights={'kernel': np.ones(3)},
+                            qcs=[build_nbits_qc(w_attr={'kernel': (8, True)})])
+        fw_info_mock.get_kernel_op_attributes = lambda nt: ['weight'] if nt is DummyLayerWKernel else ['kernel']
+
+        with pytest.raises(NotImplementedError, match='Node .*aaa with kernel cannot be used as '
+                                                      'activation for VirtualActivationWeightsNode'):
+            VirtualActivationWeightsNode(a_node, w_node, fw_info_mock)
+
diff --git a/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py b/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py
@@ -32,7 +32,7 @@
     RUTarget
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_calculator import \
     Utilization, ResourceUtilizationCalculator, TargetInclusionCriterion, BitwidthMode
-from tests_pytest._test_util.graph_builder_utils import build_node, build_qc, full_attr_name
+from tests_pytest._test_util.graph_builder_utils import build_node, full_attr_name, build_nbits_qc as build_qc
 
 BM = BitwidthMode
 TIC = TargetInclusionCriterion
@@ -844,7 +844,7 @@ def test_compute_w_utilization_non_custom(self, prepare_compute_w_util):
 
     def test_compute_w_utilization_no_targets(self, graph_mock, fw_impl_mock, fw_info_mock):
         graph_mock.nodes = [
-            build_node('n1', qcs=build_qc()),
+            build_node('n1', qcs=[build_qc()]),
             build_node('n2', canonical_weights={'foo': np.ones((5,))}, qcs=[build_qc(w_attr={'foo': (8, True)})])
         ]
         ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock)
diff --git a/tests_pytest/common_tests/unit_tests/test_model_collector.py b/tests_pytest/common_tests/unit_tests/test_model_collector.py
@@ -23,7 +23,7 @@
 from model_compression_toolkit.core.common.graph.edge import Edge
 from model_compression_toolkit.core.common.hessian import HessianInfoService
 from model_compression_toolkit.core.common.model_collector import create_stats_collector_for_node, create_tensor2node, ModelCollector
-from tests_pytest._test_util.graph_builder_utils import build_node, DummyLayer, build_qc
+from tests_pytest._test_util.graph_builder_utils import build_node, DummyLayer, build_nbits_qc as build_qc
 
 
 @pytest.fixture
