Extended test for strip

Author: ljaljushkin

tests/torch/quantization/test_strip.py

@@ -9,11 +9,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Tuple
+from typing import Any, Tuple
 
 import numpy as np
 import pytest
 import torch
+from torch import nn
 from torch.quantization.fake_quantize import FakeQuantize
 
 import nncf
@@ -26,6 +27,8 @@
 from nncf.parameters import StripFormat
 from nncf.torch.graph.transformations.commands import ExtraCompressionModuleType
 from nncf.torch.quantization.layers import AsymmetricQuantizer
+from nncf.torch.quantization.layers import BaseQuantizer
+from nncf.torch.quantization.layers import BaseWeightsDecompressor
 from nncf.torch.quantization.layers import PTQuantizerSpec
 from nncf.torch.quantization.layers import SymmetricQuantizer
 from nncf.torch.quantization.strip import convert_to_torch_fakequantizer
@@ -330,6 +333,30 @@ def test_nncf_strip_api(strip_type, do_copy):
     assert isinstance(strip_model.nncf.external_quantizers["/nncf_model_input_0|OUTPUT"], FakeQuantize)
 
 
+def check_compression_modules(
+    model_: nn.Module,
+    expected_module_type: ExtraCompressionModuleType,
+    not_expected_module_type: ExtraCompressionModuleType,
+    expected_class: Any,
+) -> None:
+    """
+    Checks if the given model has the expected compression module registered and not the unexpected one.
+    Also verifies that the compression module is of the expected class type.
+
+    :param model_: The model to be checked, which should have an 'nncf' attribute with compression module methods.
+    :param expected_module_type: The type of the compression module that is expected to be registered.
+    :param not_expected_module_type: The type of the compression module that is not expected to be registered.
+    :param expected_class: The class type that the expected compression module should be an instance of.
+    """
+    print(model_)
+    assert model_.nncf.is_compression_module_registered(expected_module_type)
+    assert not model_.nncf.is_compression_module_registered(not_expected_module_type)
+    compression_modules_dict = model_.nncf.get_compression_modules_by_type(expected_module_type)
+    assert len(compression_modules_dict) == 1
+    compression_module = next(iter(compression_modules_dict.values()))
+    assert isinstance(compression_module, expected_class)
+
+
 @pytest.mark.parametrize(
     ("mode", "torch_dtype", "atol"),
     (
@@ -358,11 +385,23 @@ def test_nncf_strip_lora_model(mode, torch_dtype, atol):
     if mode in [CompressWeightsMode.INT4_SYM, CompressWeightsMode.INT4_ASYM]:
         compression_kwargs.update(dict(ratio=1, group_size=4, all_layers=True))
     compressed_model = nncf.compress_weights(model, **compression_kwargs)
+    check_compression_modules(
+        compressed_model,
+        expected_module_type=ExtraCompressionModuleType.EXTERNAL_QUANTIZER,
+        not_expected_module_type=ExtraCompressionModuleType.EXTERNAL_OP,
+        expected_class=BaseQuantizer,
+    )
 
     with torch.no_grad():
         compressed_output = compressed_model(example)
 
         strip_compressed_model = nncf.strip(compressed_model, do_copy=True, strip_format=StripFormat.DQ)
+        check_compression_modules(
+            strip_compressed_model,
+            expected_module_type=ExtraCompressionModuleType.EXTERNAL_OP,
+            not_expected_module_type=ExtraCompressionModuleType.EXTERNAL_QUANTIZER,
+            expected_class=BaseWeightsDecompressor,
+        )
         stripped_output = strip_compressed_model(example)
 
         assert torch.allclose(compressed_output, stripped_output, atol=atol)