Apply activation quantization parameters selection(2nd PR internal review) (#15)

Fixed for PR#1466 review comments.

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
@@ -32,6 +32,7 @@
     OpQuantizationConfig
 
 if TYPE_CHECKING:
+    from model_compression_toolkit.core.common import BaseNode
     from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
 
 ##########################################
@@ -199,15 +200,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
 

tests_pytest/pytorch_tests/unit_tests/core/common/quantization/quantization_params_generation/test_calculate_quantization_params.py

@@ -23,17 +23,22 @@
     NodeActivationQuantizationConfig, NodeWeightsQuantizationConfig
 from model_compression_toolkit.target_platform_capabilities import OpQuantizationConfig
 from model_compression_toolkit.core import QuantizationConfig, QuantizationErrorMethod
+from model_compression_toolkit.core.common.hessian.hessian_info_service import HessianInfoService
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2pytorch import \
     AttachTpcToPytorch
 import model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema as schema
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import Signedness, \
     AttributeQuantizationConfig
-from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.core.pytorch.default_framework_info import PyTorchInfo
+from model_compression_toolkit.core.common.framework_info import set_fw_info, get_fw_info
+
 from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
 from model_compression_toolkit.core.common.collectors.statistics_collector import StatsCollector
 from model_compression_toolkit.target_platform_capabilities.constants import KERNEL_ATTR, WEIGHTS_N_BITS
 from mct_quantizers import QuantizationMethod
 
+from model_compression_toolkit.core.common.framework_info import ChannelAxisMapping
+
 class TestCalculateQuantizationParams:
     def get_op_qco(self):
         # define a default quantization config for all non-specified weights attributes.
@@ -145,12 +150,11 @@ def _create_node_weights_op_cfg(self,
 
     def get_test_graph(self, qem: QuantizationErrorMethod):
         float_model = self.get_float_model()
-        fw_info = DEFAULT_PYTORCH_INFO
+        set_fw_info(PyTorchInfo)
 
         fw_impl = PytorchImplementation()
         graph = fw_impl.model_reader(float_model,
                                      self.representative_data_gen)
-        graph.set_fw_info(fw_info)
 
         quantization_config = QuantizationConfig(weights_error_method=qem)
 
@@ -165,24 +169,23 @@ def get_test_graph(self, qem: QuantizationErrorMethod):
 
         graph.node_to_out_stats_collector = dict()
         for id, n in enumerate(graph.nodes):
-            n.prior_info = fw_impl.get_node_prior_info(node=n, fw_info=fw_info, graph=graph)
+            n.prior_info = fw_impl.get_node_prior_info(node=n, graph=graph)
             n.candidates_quantization_cfg = []
             candidate_qc_a = CandidateNodeQuantizationConfig(
                 activation_quantization_cfg=NodeActivationQuantizationConfig(qc=quantization_config, op_cfg=op_cfg,
                                                                              activation_quantization_fn=None,
                                                                              activation_quantization_params_fn=None),
                 weights_quantization_cfg=NodeWeightsQuantizationConfig(qc=quantization_config, op_cfg=op_cfg,
-                                                                       weights_channels_axis=(0, 1),
+                                                                       weights_channels_axis=ChannelAxisMapping(0, 1),
                                                                        node_attrs_list=['weight', 'bias'])
             )
             if n.name in ['conv3']:
                 candidate_qc_a.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.FLN_QUANT
-                candidate_qc_a.activation_quantization_cfg.activation_n_bits = 16 # set 16bit for FLN node for test.
             else:
                 candidate_qc_a.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.QUANT
             n.candidates_quantization_cfg.append(candidate_qc_a)
 
-            graph.node_to_out_stats_collector[n] = StatsCollector(init_min_value=0.0, init_max_value=1.0, out_channel_axis=fw_info.out_channel_axis_mapping.get(n.type))
+            graph.node_to_out_stats_collector[n] = StatsCollector(init_min_value=0.0, init_max_value=1.0, out_channel_axis=get_fw_info().out_channel_axis_mapping.get(n.type))
             graph.node_to_out_stats_collector[n].hc._n_bins = 3
             if n.name in ['conv1']:
                 graph.node_to_out_stats_collector[n].hc._bins = np.array([0.4, 0.8, 1.2])