Adding saturation fix parameters to compression state for PT (openvinotoolkit#1128)

SKholkin · web-flow · commit da72001a9e0f · 2022-04-04T15:40:22.000+03:00
* Trying to debug autoq tests failure

* First working iteration

* New compression state loading and backward compatibility

* Rebased

* Test NotImplementedError deleted

* Little clean up

* Compression State versioning

* Comparison operators

* nas test fix

* Composite compression loading

* Enum PT compression state versioning

* Remove accidentally added files

* Changed version to local quantizaer builder one

* Fixed pylint

* Fixes
diff --git a/nncf/torch/checkpoint_loading.py b/nncf/torch/checkpoint_loading.py
@@ -18,7 +18,6 @@
 import torch
 
 from nncf.common.utils.logger import logger as nncf_logger
-from nncf.torch.utils import maybe_convert_legacy_names_in_model_state
 
 
 def load_state(model: torch.nn.Module, state_dict_to_load: dict, is_resume: bool = False,
@@ -45,6 +44,7 @@ def load_state(model: torch.nn.Module, state_dict_to_load: dict, is_resume: bool
 
     model_state_dict = model.state_dict()
 
+    from nncf.torch.utils import maybe_convert_legacy_names_in_model_state
     maybe_convert_legacy_names_in_model_state(state_dict_to_load)
     key_matcher = KeyMatcher(is_resume, state_dict_to_load, model_state_dict, keys_to_ignore)
     new_dict = key_matcher.run()
diff --git a/nncf/torch/quantization/algo.py b/nncf/torch/quantization/algo.py
diff --git a/nncf/torch/quantization/layers.py b/nncf/torch/quantization/layers.py
@@ -11,7 +11,7 @@
  limitations under the License.
 """
 from enum import Enum
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Optional, Any
 
 import numpy as np
 import torch
@@ -24,10 +24,15 @@
 from nncf.torch.checkpoint_loading import OPTIONAL_PARAMETERS_REGISTRY
 from nncf.common.utils.debug import is_debug
 from nncf.torch.functions import clamp
+from nncf.common.graph import NNCFNodeName
 from nncf.common.utils.logger import logger as nncf_logger
 from nncf.common.quantization.structs import QuantizationMode, QuantizerConfig, QuantizerSpec
 from nncf.common.quantization.quantizers import calculate_symmetric_level_ranges
 from nncf.common.quantization.quantizers import calculate_asymmetric_level_ranges
+from nncf.common.quantization.quantizer_setup import QuantizerSetupBase
+from nncf.common.quantization.quantizer_setup import QuantizationPointId
+from nncf.torch.graph.transformations.commands import TargetType
+from nncf.torch.graph.transformations.commands import PTTargetPoint
 from nncf.torch.quantization.quantize_functions import symmetric_quantize, asymmetric_quantize, \
     ExportQuantizeToFakeQuantize, get_scale_zp_from_input_low_input_high, ExportQuantizeToONNXQuantDequant, TuneRange
 from nncf.torch.layer_utils import COMPRESSION_MODULES, CompressionParameter
@@ -52,7 +57,21 @@ def from_str(config_value: str) -> 'HWConfigType':
         raise RuntimeError("Unknown quantizer ONNX export mode string")
 
 
+class PTQSpecStateNames:
+    NUM_BITS = 'num_bits'
+    MODE = 'mode'
+    SIGNED_TO_FORCE = 'signedness_to_force'
+    NARROW_RANGE = 'narrow_range'
+    HALF_RANGE = 'half_range'
+    SCALE_SHAPE = 'scale_shape'
+    LOGARITHM_SCALE = 'logarithm_scale'
+    IS_QUANTIZED_ON_EXPORT = 'is_quantized_on_export'
+    COMPRESSION_LR_MULTIPLIER = 'compression_lr_multiplier'
+
+
 class PTQuantizerSpec(QuantizerSpec):
+    _state_names = PTQSpecStateNames
+
     def __init__(self, num_bits: int,
                  mode: QuantizationMode,
                  signedness_to_force: Optional[bool],
@@ -70,6 +89,7 @@ def __init__(self, num_bits: int,
             activation quantizers.
         """
         super().__init__(num_bits, mode, signedness_to_force, narrow_range, half_range)
+        self.per_channel = scale_shape != [1]
         self.scale_shape = scale_shape
         self.logarithm_scale = logarithm_scale
         self.compression_lr_multiplier = compression_lr_multiplier
@@ -90,6 +110,155 @@ def from_config(cls, qconfig: QuantizerConfig, narrow_range: bool,
                    is_quantized_on_export,
                    compression_lr_multiplier)
 
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    @classmethod
+    def from_state(cls, state: Dict[str, Any]) -> 'PTQuantizationPoint':
+        """
+        Creates the object from its state.
+
+        :param state: Output of `get_state()` method.
+        """
+        kwargs = {
+            cls._state_names.NUM_BITS: state['num_bits'],
+            cls._state_names.MODE: state['mode'],
+            cls._state_names.SIGNED_TO_FORCE: state['signedness_to_force'],
+            cls._state_names.NARROW_RANGE: state['narrow_range'],
+            cls._state_names.HALF_RANGE: state['half_range'],
+            cls._state_names.SCALE_SHAPE: state['scale_shape'],
+            cls._state_names.LOGARITHM_SCALE: state['logarithm_scale'],
+            cls._state_names.IS_QUANTIZED_ON_EXPORT: state['is_quantized_on_export'],
+            cls._state_names.COMPRESSION_LR_MULTIPLIER: state['compression_lr_multiplier']
+        }
+        return cls(**kwargs)
+
+    def get_state(self):
+        return {self._state_names.NUM_BITS: self.num_bits,
+                self._state_names.MODE: self.mode,
+                self._state_names.SIGNED_TO_FORCE: self.signedness_to_force,
+                self._state_names.NARROW_RANGE: self.narrow_range,
+                self._state_names.HALF_RANGE: self.half_range,
+                self._state_names.SCALE_SHAPE: self.scale_shape,
+                self._state_names.LOGARITHM_SCALE: self.logarithm_scale,
+                self._state_names.IS_QUANTIZED_ON_EXPORT: self.is_quantized_on_export,
+                self._state_names.COMPRESSION_LR_MULTIPLIER: self.compression_lr_multiplier}
+
+
+class PTQPointStateNames:
+    QSPEC = 'qspec'
+    TARGET_POINT = 'target_point'
+    NAMES_OF_QUANTIZED_OPS = 'directly_quantized_operator_node_names'
+
+
+class PTQuantizationPoint:
+    _state_names = PTQPointStateNames
+
+    def __init__(self, qspec: PTQuantizerSpec, target_point: PTTargetPoint,
+                 directly_quantized_operator_node_names: List[NNCFNodeName]):
+        self.qspec = qspec
+        self.target_point = target_point
+        self.directly_quantized_operator_node_names = directly_quantized_operator_node_names
+
+    def is_activation_quantization_point(self) -> bool:
+        return not self.is_weight_quantization_point()
+
+    def is_weight_quantization_point(self) -> bool:
+        return self.target_point.target_type == TargetType.OPERATION_WITH_WEIGHTS
+
+    def __str__(self):
+        return str(self.target_point) + ' ' + str(self.qspec)
+
+    def get_state(self) -> Dict[str, Any]:
+        """
+        Returns a dictionary with Python data structures (dict, list, tuple, str, int, float, True, False, None) that
+        represents state of the object.
+
+        :return: state of the object
+        """
+        return {
+            self._state_names.TARGET_POINT: self.target_point.get_state(),
+            self._state_names.QSPEC: self.qspec.get_state(),
+            self._state_names.NAMES_OF_QUANTIZED_OPS: self.directly_quantized_operator_node_names
+        }
+
+    @classmethod
+    def from_state(cls, state: Dict[str, Any]) -> 'PTQuantizationPoint':
+        """
+        Creates the object from its state.
+
+        :param state: Output of `get_state()` method.
+        """
+        kwargs = {
+            cls._state_names.TARGET_POINT: PTTargetPoint.from_state(state[cls._state_names.TARGET_POINT]),
+            cls._state_names.QSPEC: PTQuantizerSpec.from_state(state[cls._state_names.QSPEC]),
+            cls._state_names.NAMES_OF_QUANTIZED_OPS: state[cls._state_names.NAMES_OF_QUANTIZED_OPS]
+        }
+        return cls(**kwargs)
+
+
+class PTQSetupStateNames:
+    SHARED_INPUT_OPERATION_SET_GROUPS = 'shared_input_operation_set_groups'
+    UNIFIED_SCALE_GROUPS = 'unified_scale_groups'
+    QUANTIZATION_POINTS = 'quantization_points'
+
+
+class PTQuantizerSetup(QuantizerSetupBase):
+    _state_names = PTQSetupStateNames
+
+    def __init__(self, unified_scale_groups, shared_input_operation_set_groups):
+        super().__init__()
+        self.unified_scale_groups = unified_scale_groups
+        self.shared_input_operation_set_groups = shared_input_operation_set_groups
+        self.quantization_points = {}  # type: Dict[QuantizationPointId, PTQuantizationPoint]
+
+    @classmethod
+    def from_state(cls, state: Dict) -> 'PTQuantizerSetup':
+        """
+        Creates the object from its state.
+
+        :param state: Output of `get_state()` method.
+        """
+
+        def decode_qp(pair):
+            str_qp_id, qp_state = pair
+            return int(str_qp_id), PTQuantizationPoint.from_state(qp_state)
+
+        def list2set(pair):
+            str_idx, qp_id_list = pair
+            return int(str_idx), set(qp_id_list)
+
+        unified_scale_groups = dict(map(list2set, state[cls._state_names.UNIFIED_SCALE_GROUPS].items()))
+        shared_input_operation_set_groups_state = state[cls._state_names.SHARED_INPUT_OPERATION_SET_GROUPS]
+        setup = PTQuantizerSetup(unified_scale_groups, shared_input_operation_set_groups_state)
+        setup.quantization_points = dict(map(decode_qp, state[cls._state_names.QUANTIZATION_POINTS].items()))
+        setup.shared_input_operation_set_groups = dict(map(list2set, shared_input_operation_set_groups_state.items()))
+        return setup
+
+    def get_state(self):
+        """
+        Returns a dictionary with Python data structures (dict, list, tuple, str, int, float, True, False, None) that
+        represents state of the object.
+
+        :return: state of the object
+        """
+
+        def set2list(pair):
+            i, qp_id_set = pair
+            return i, list(qp_id_set)
+
+        quantization_points_state = {qp_id: qp.get_state() for qp_id, qp in self.quantization_points.items()}
+        unified_scale_groups_state = dict(map(set2list, self.unified_scale_groups.items()))
+        shared_input_operation_set_groups_state = dict(map(set2list, self.shared_input_operation_set_groups.items()))
+        return {
+            self._state_names.QUANTIZATION_POINTS: quantization_points_state,
+            self._state_names.UNIFIED_SCALE_GROUPS: unified_scale_groups_state,
+            self._state_names.SHARED_INPUT_OPERATION_SET_GROUPS: shared_input_operation_set_groups_state,
+        }
+
+    def add_quantization_point(self, qp_id: QuantizationPointId, qp: PTQuantizationPoint):
+        self.quantization_points[qp_id] = qp
+
 
 class BaseQuantizer(nn.Module):
     # pylint:disable=too-many-public-methods
@@ -636,7 +805,6 @@ def get_quantizer_config(self) -> QuantizerConfig:
                                per_channel=self.per_channel)
 
 
-
 def get_per_channel_scale_shape(input_shape, is_weights, channel_idx: int = None):
     scale_shape = [1 for _ in input_shape]
     if channel_idx is None:
diff --git a/nncf/torch/quantization/precision_init/adjacent_quantizers.py b/nncf/torch/quantization/precision_init/adjacent_quantizers.py
@@ -80,7 +80,7 @@ def parse_from_quantizer_setup(self, all_quantizations: Dict[QuantizerId, BaseQu
                 resulting_tuple = (quant_id, quantizer_module)
                 if qp.is_weight_quantization_point():
                     wt_quant_tuples.append(resulting_tuple)
-                    weight_quantized_module_node_name = qp.insertion_point.target_node_name
+                    weight_quantized_module_node_name = qp.target_point.target_node_name
                     module_scope_per_weight_qp_id[weight_quantized_module_node_name] = qp_id
                 elif qp.is_activation_quantization_point():
                     act_quant_tuples.append(resulting_tuple)
diff --git a/nncf/torch/utils.py b/nncf/torch/utils.py
@@ -346,6 +346,7 @@ def maybe_convert_legacy_names_in_model_state(state_dict_to_load: Dict[str, Any]
     for old_name, new_name in LEGACY_VS_NEW_BN_MAP.items():
         rename_legacy_names_in_state_dict(state_dict_to_load, legacy_names[old_name], old_name, new_name)
 
+
 def maybe_convert_legacy_names_in_compress_state(compression_state: Dict[str, Any]) -> None:
     """
     Convert legacy layer names in compression state in case such names exist.
@@ -359,29 +360,32 @@ def maybe_convert_legacy_names_in_compress_state(compression_state: Dict[str, An
     if not controller_state or 'quantization' not in controller_state:
         return
 
-    qips = controller_state['quantization']['quantizer_setup']['quantization_points']
-
-    detected_legacy_names = {
-        'BatchNorm1d': False,
-        'BatchNorm2d': False,
-        'BatchNorm3d': False,
-        'NNCFBatchNorm': False,
-    }
+    from nncf.torch.quantization.algo import QUANTIZER_BUILDER_STATE_VERSION_SAVE_NAME
+    if not controller_state['quantization'].get(QUANTIZER_BUILDER_STATE_VERSION_SAVE_NAME):
+        qips = controller_state['quantization']['quantizer_setup']['quantization_points']
+
+        detected_legacy_names = {
+            'BatchNorm1d': False,
+            'BatchNorm2d': False,
+            'BatchNorm3d': False,
+            'NNCFBatchNorm': False,
+        }
+
+        for point in qips.values():
+            name = point['qip']['target_node_name']
+            for old_name, new_name in LEGACY_VS_NEW_BN_MAP.items():
+                if old_name in name and not new_name in name:
+                    detected_legacy_names[old_name] = True
+                    point['qip']['target_node_name'] = name.replace(old_name, new_name)
+                    break
+
+        for old_name, was_detected in detected_legacy_names.items():
+            if was_detected:
+                new_name = LEGACY_VS_NEW_BN_MAP[old_name]
+                warning_deprecated('Legacy Batch Norm layer names was detected in quantization setup target'
+                                   ' point names. All occurrences of `{}` in nodes names was replaced by'
+                                   ' `{}`'.format(old_name, new_name))
 
-    for point in qips.values():
-        name = point['qip']['target_node_name']
-        for old_name, new_name in LEGACY_VS_NEW_BN_MAP.items():
-            if old_name in name and not new_name in name:
-                detected_legacy_names[old_name] = True
-                point['qip']['target_node_name'] = name.replace(old_name, new_name)
-                break
-
-    for old_name, was_detected in  detected_legacy_names.items():
-        if was_detected:
-            new_name = LEGACY_VS_NEW_BN_MAP[old_name]
-            warning_deprecated('Legacy Batch Norm layer names was detected in quantization setup target point names. '
-                               'All occurrences of `{}` in nodes names was replaced by `{}`'.format(old_name,
-                                                                                                    new_name))
 
 def get_model_device(model: torch.nn.Module) -> torch.device:
     try:
diff --git a/tests/torch/nas/test_all_elasticity.py b/tests/torch/nas/test_all_elasticity.py
@@ -366,7 +366,8 @@ def test_multi_elasticity_state():
     prepare_train_algo_for_resume(training_ctrl)
     compression_state = training_ctrl.get_compression_state()
 
-    assert compression_state == REF_COMPRESSION_STATE_FOR_TWO_CONV
+    assert compression_state['ctrl_state'] == REF_COMPRESSION_STATE_FOR_TWO_CONV['ctrl_state']
+    assert compression_state['builder_state'] == REF_COMPRESSION_STATE_FOR_TWO_CONV['builder_state']
 
 
 def test_can_restore_from_state():
diff --git a/tests/torch/quantization/test_hawq_precision_init.py b/tests/torch/quantization/test_hawq_precision_init.py
@@ -780,10 +780,10 @@ def test_compression_ratio(desc, mocker):
     config = desc.create_config()
     register_bn_adaptation_init_args(config)
     from nncf.torch.quantization.algo import QuantizationBuilder
-    get_qsetyp_spy = mocker.spy(QuantizationBuilder, '_get_quantizer_setup')
+    get_single_config_quantizer_setup_spy = mocker.spy(QuantizationBuilder, '_get_single_config_quantizer_setup')
     model, ctrl = create_compressed_model_and_algo_for_test(ConvLinear(), config)
 
-    quantizer_setup = get_qsetyp_spy.spy_return
+    quantizer_setup = get_single_config_quantizer_setup_spy.spy_return
     weight_qp_id_per_activation_qp_id = ctrl.groups_of_adjacent_quantizers.weight_qp_id_per_activation_qp_id
     flops_per_module = model.get_flops_per_module()
     ratio_calculator = CompressionRatioCalculator(flops_per_module, quantizer_setup, weight_qp_id_per_activation_qp_id)
diff --git a/tests/torch/test_backward_compat.py b/tests/torch/test_backward_compat.py
diff --git a/tests/torch/test_compressed_graph.py b/tests/torch/test_compressed_graph.py
