set reuse in exporter

Author: eladc-git

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -73,6 +73,11 @@ def export(self) -> None:
             """
             for layer in self.model.children():
                 self.is_layer_exportable_fn(layer)
+                # Set reuse for weight quantizers if quantizer is reused
+                if isinstance(layer, PytorchQuantizationWrapper):
+                    for _, quantizer in layer.weights_quantizers.items():
+                        if quantizer.reuse:
+                            quantizer.enable_reuse_quantizer()
 
             # Set forward that is used during onnx export.
             # If _use_onnx_custom_quantizer_ops is set to True, the quantizer forward function will use
@@ -116,6 +121,13 @@ def export(self) -> None:
                                   dynamic_axes={'input': {0: 'batch_size'},
                                                 'output': {0: 'batch_size'}})
 
+            for layer in self.model.children():
+                # Set disable for reuse for weight quantizers if quantizer was reused
+                if isinstance(layer, PytorchQuantizationWrapper):
+                    for _, quantizer in layer.weights_quantizers.items():
+                        if quantizer.reuse:
+                            quantizer.disable_reuse_quantizer()
+
         def _enable_onnx_custom_ops_export(self):
             """
             Enable the custom implementation forward in quantizers, so it is exported

model_compression_toolkit/exporter/model_wrapper/pytorch/builder/fully_quantized_model_builder.py

@@ -48,7 +48,7 @@ def fully_quantized_wrapper(node: common.BaseNode,
             # Set reuse for weight quantizers if node is reused
             for _, quantizer in weight_quantizers.items():
                 if node.reuse_group:
-                    quantizer.enable_reuse_quantizer()
+                    quantizer.reuse = True
             # for positional weights we need to extract the weight's value.
             weights_values = {attr: fw_impl.to_tensor(node.get_weights_by_keys(attr))
                               for attr in weight_quantizers if isinstance(attr, int)}