modify for a PR review comment(SonySemiconductorSolutions#1466 (comment))

Author: gouda-youichi

model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
+from __future__ import annotations
 
 from typing import Callable, Any, List, Tuple, Union, Dict, TYPE_CHECKING
 from enum import Enum, auto
@@ -31,6 +31,9 @@
     AttributeQuantizationConfig, \
     OpQuantizationConfig
 
+if TYPE_CHECKING:
+    from model_compression_toolkit.core.common import BaseNode
+
 if TYPE_CHECKING:
     from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
 
@@ -199,15 +202,17 @@ def set_activation_quantization_params_fn(self, activation_quantization_params_f
         self.activation_quantization_params_fn = activation_quantization_params_fn
 
     def set_activation_quantization_param(self,
-                                          activation_params: dict):
+                                          activation_params: dict,
+                                          node: BaseNode):
         """
          Set a quantization parameter for the node's activation.
 
         Args:
-            activation_params: Dictionary that contains weight quantization params.
+            activation_params: Dictionary that contains activation quantization params.
+            node: node in a graph that represents the model.
 
         """
-        assert self.quant_mode == ActivationQuantizationMode.QUANT or self.quant_mode == ActivationQuantizationMode.FLN_QUANT
+        assert node.is_activation_quantization_enabled() or node.is_fln_quantization()
         for param_name, param_value in activation_params.items():
             self.activation_quantization_params[param_name] = param_value
 

model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -136,4 +136,4 @@ def calculate_quantization_params(graph: Graph,
                     nodes_prior_info=n.prior_info,
                     out_stats_container=graph.get_out_stats_collector(n))
                 # Create a NodeQuantizationConfig containing all quantization params and attach it to the node
-                candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_params)
+                candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_params, n)

model_compression_toolkit/core/common/statistics_correction/apply_second_moment_correction_to_graph.py

@@ -56,7 +56,7 @@ def _collect_and_assign_act_threshold(graph: Graph,
                 activation_quant_cfg=n.final_activation_quantization_cfg,
                 nodes_prior_info=n.prior_info,
                 out_stats_container=graph.get_out_stats_collector(n))
-            n.final_activation_quantization_cfg.set_activation_quantization_param(activation_params)
+            n.final_activation_quantization_cfg.set_activation_quantization_param(activation_params, n)
 
 
 def quantized_model_builder_for_second_moment_correction(graph: common.Graph,

model_compression_toolkit/core/common/substitutions/shift_negative_activation.py

@@ -465,7 +465,7 @@ def shift_negative_function(graph: Graph,
                                                                              add_node_qco[op_qc_idx])
 
         candidate_qc.activation_quantization_cfg.set_activation_quantization_param({THRESHOLD: activation_threshold,
-                                                                                    SIGNED: False})
+                                                                                    SIGNED: False}, add_node)
 
         candidate_qc.activation_quantization_cfg.activation_n_bits = original_non_linear_activation_nbits
 
@@ -482,7 +482,7 @@ def shift_negative_function(graph: Graph,
 
         assert activation_param.get(SIGNED) == False
         for candidate_qc in non_linear_node.candidates_quantization_cfg:
-            candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_param)
+            candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_param, non_linear_node)
 
     return graph