Fix FQ exporter for 16bit act quant and constant quantization (#1463)

Remove _replace_wrapped_with_unwrapped from FQ exporter as it's not needed.
Remove PytorchActivationQuantizationHolder with num_bits > 8, because ONNX doesn't support it.

---------

Co-authored-by: Ofir Gordon <ofirgo@sony.com>

Author: elad-c

model_compression_toolkit/exporter/model_exporter/pytorch/base_pytorch_exporter.py

@@ -56,7 +56,7 @@ def _search(m):
                 f"Only the first one was assigned to 'model.metadata'.")
 
 
-def _set_quantized_weights_in_wrapper(layer:PytorchQuantizationWrapper):
+def _set_quantized_weights_in_wrapper(layer: PytorchQuantizationWrapper):
     """
        Sets the quantized weights in the provided PytorchQuantizationWrapper layer.
        Replaces the original weights in the layer with the quantized weights.
@@ -124,7 +124,7 @@ def __init__(self,
         self.model = copy.deepcopy(self.model)
         self.repr_dataset = repr_dataset
 
-    def _substitute_fully_quantized_model(self):
+    def _substitute_fully_quantized_model(self, replace_wrapped=True):
         """
         Substitution for pytorch "fully-quantized" models. It first uses the weight quantizers
         in PytorchQuantizationWrapper layers to quantize the weights and set them in the layer.
@@ -136,8 +136,9 @@ def _substitute_fully_quantized_model(self):
             if isinstance(layer, PytorchQuantizationWrapper):
                 _set_quantized_weights_in_wrapper(layer)
 
-        # Replace PytorchQuantizationWrapper layers with their internal layers
-        self._replace_wrapped_with_unwrapped()
+        if replace_wrapped:
+            # Replace PytorchQuantizationWrapper layers with their internal layers
+            self._replace_wrapped_with_unwrapped()
 
     def _replace_wrapped_with_unwrapped(self):
         """
@@ -148,3 +149,4 @@ def _replace_wrapped_with_unwrapped(self):
         for name, module in self.model.named_children():
             if isinstance(module, PytorchQuantizationWrapper):
                 setattr(self.model, name, module.layer)
+

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -73,6 +73,15 @@ def export(self, output_names=None) -> None:
             Returns:
                 Fake-quant PyTorch model.
             """
+            # List all activation quantization holders with num_bits>8 and replace them with Identity, because
+            # ONNX doesn't support quantization of more than 8 bits for torch.fake_quantize_per_tensor_affine.
+            act_holder_list = [n for n, m in self.model.named_modules()
+                               if isinstance(m, PytorchActivationQuantizationHolder) and
+                               m.activation_holder_quantizer.num_bits > 8]
+            for act_holder in act_holder_list:  # pragma: no cover
+                delattr(self.model, act_holder)
+                setattr(self.model, act_holder, torch.nn.Identity())
+
             for layer in self.model.children():
                 self.is_layer_exportable_fn(layer)
                 # Set reuse for weight quantizers if quantizer is reused
@@ -89,7 +98,7 @@ def export(self, output_names=None) -> None:
             if self._use_onnx_custom_quantizer_ops:
                 self._enable_onnx_custom_ops_export()
             else:
-                self._substitute_fully_quantized_model()
+                self._substitute_fully_quantized_model(replace_wrapped=False)
 
             if self._use_onnx_custom_quantizer_ops:
                 Logger.info(f"Exporting onnx model with MCTQ quantizers: {self.save_model_path}")
@@ -166,6 +175,6 @@ def _enable_onnx_custom_ops_export(self):
                         wq.enable_custom_impl()
 
 else:
-    def FakelyQuantONNXPyTorchExporter(*args, **kwargs):
+    def FakelyQuantONNXPyTorchExporter(*args, **kwargs):  # pragma: no cover
         Logger.critical("ONNX must be installed to use 'FakelyQuantONNXPyTorchExporter'. "
-                        "The 'onnx' package is missing.")  # pragma: no cover
+                        "The 'onnx' package is missing.")

tests/pytorch_tests/function_tests/test_export_pytorch_fully_quantized_model.py

@@ -98,7 +98,7 @@ def to_numpy(tensor):
             return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
         # get onnx conv weight
-        conv_weight_name = 'conv1.weight'
+        conv_weight_name = 'conv1.layer.weight'
         from onnx import numpy_helper
         INTIALIZERS = self.exported_model_onnx.graph.initializer
         Weight = []