mct-nightly 1.7.1.31122022.post351__py3-none-any.whl → 1.8.0.1042023.post423__py3-none-any.whl

model_compression_toolkit/core/tpc_models/imx500_tpc/v1/tpc_keras.py

@@ -0,0 +1,101 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import tensorflow as tf
+from packaging import version
+
+if version.parse(tf.__version__) < version.parse("2.6"):
+    from tensorflow.keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
+        MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
+        Conv2DTranspose
+else:
+    from keras.layers import Conv2D, DepthwiseConv2D, Dense, Reshape, ZeroPadding2D, Dropout, \
+        MaxPooling2D, Activation, ReLU, Add, Subtract, Multiply, PReLU, Flatten, Cropping2D, LeakyReLU, Permute, \
+        Conv2DTranspose
+
+from model_compression_toolkit.core.tpc_models.imx500_tpc.v1.tp_model import get_tp_model
+import model_compression_toolkit as mct
+from model_compression_toolkit.core.tpc_models.imx500_tpc.v1 import __version__ as TPC_VERSION
+
+tp = mct.target_platform
+
+
+def get_keras_tpc() -> tp.TargetPlatformCapabilities:
+    """
+    get a Keras TargetPlatformCapabilities object with default operation sets to layers mapping.
+    Returns: a Keras TargetPlatformCapabilities object for the given TargetPlatformModel.
+    """
+    imx500_tpc_tp_model = get_tp_model()
+    return generate_keras_tpc(name='imx500_tpc_keras_tpc', tp_model=imx500_tpc_tp_model)
+
+
+def generate_keras_tpc(name: str, tp_model: tp.TargetPlatformModel):
+    """
+    Generates a TargetPlatformCapabilities object with default operation sets to layers mapping.
+
+    Args:
+        name: Name of the TargetPlatformCapabilities.
+        tp_model: TargetPlatformModel object.
+
+    Returns: a TargetPlatformCapabilities object for the given TargetPlatformModel.
+    """
+
+    keras_tpc = tp.TargetPlatformCapabilities(tp_model, name=name, version=TPC_VERSION)
+
+    with keras_tpc:
+        tp.OperationsSetToLayers("NoQuantization", [Reshape,
+                                                    tf.reshape,
+                                                    Permute,
+                                                    tf.transpose,
+                                                    Flatten,
+                                                    Cropping2D,
+                                                    ZeroPadding2D,
+                                                    Dropout,
+                                                    MaxPooling2D,
+                                                    tf.split,
+                                                    tf.quantization.fake_quant_with_min_max_vars,
+                                                    tf.math.argmax,
+                                                    tf.shape,
+                                                    tf.math.equal,
+                                                    tf.gather,
+                                                    tf.cast,
+                                                    tf.compat.v1.gather,
+                                                    tf.nn.top_k,
+                                                    tf.__operators__.getitem,
+                                                    tf.compat.v1.shape])
+
+        tp.OperationsSetToLayers("Conv", [Conv2D,
+                                          DepthwiseConv2D,
+                                          Conv2DTranspose,
+                                          tf.nn.conv2d,
+                                          tf.nn.depthwise_conv2d,
+                                          tf.nn.conv2d_transpose])
+        tp.OperationsSetToLayers("FullyConnected", [Dense])
+        tp.OperationsSetToLayers("AnyReLU", [tf.nn.relu,
+                                             tf.nn.relu6,
+                                             tf.nn.leaky_relu,
+                                             ReLU,
+                                             LeakyReLU,
+                                             tp.LayerFilterParams(Activation, activation="relu"),
+                                             tp.LayerFilterParams(Activation, activation="leaky_relu")])
+        tp.OperationsSetToLayers("Add", [tf.add, Add])
+        tp.OperationsSetToLayers("Sub", [tf.subtract, Subtract])
+        tp.OperationsSetToLayers("Mul", [tf.math.multiply, Multiply])
+        tp.OperationsSetToLayers("Div", [tf.math.divide])
+        tp.OperationsSetToLayers("PReLU", [PReLU])
+        tp.OperationsSetToLayers("Swish", [tf.nn.swish, tp.LayerFilterParams(Activation, activation="swish")])
+        tp.OperationsSetToLayers("Sigmoid", [tf.nn.sigmoid, tp.LayerFilterParams(Activation, activation="sigmoid")])
+        tp.OperationsSetToLayers("Tanh", [tf.nn.tanh, tp.LayerFilterParams(Activation, activation="tanh")])
+
+    return keras_tpc

model_compression_toolkit/core/tpc_models/imx500_tpc/v1/tpc_pytorch.py

@@ -0,0 +1,95 @@
+# Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import operator
+
+import torch
+from torch import add, sub, mul, div, flatten, reshape, split, unsqueeze, dropout, sigmoid, tanh, chunk, unbind, topk, \
+    gather, equal, transpose, permute
+from torch.nn import Conv2d, Linear, BatchNorm2d, ConvTranspose2d
+from torch.nn import Dropout, Flatten, Hardtanh
+from torch.nn import ReLU, ReLU6, PReLU, SiLU, Sigmoid, Tanh, Hardswish, LeakyReLU
+from torch.nn.functional import relu, relu6, prelu, silu, hardtanh, hardswish, leaky_relu
+
+from model_compression_toolkit.core.tpc_models.imx500_tpc.v1.tp_model import get_tp_model
+import model_compression_toolkit as mct
+from model_compression_toolkit.core.tpc_models.imx500_tpc.v1 import __version__ as TPC_VERSION
+
+tp = mct.target_platform
+
+
+def get_pytorch_tpc() -> tp.TargetPlatformCapabilities:
+    """
+    get a Pytorch TargetPlatformCapabilities object with default operation sets to layers mapping.
+    Returns: a Pytorch TargetPlatformCapabilities object for the given TargetPlatformModel.
+    """
+    imx500_tpc_tp_model = get_tp_model()
+    return generate_pytorch_tpc(name='imx500_tpc_pytorch_tpc', tp_model=imx500_tpc_tp_model)
+
+
+def generate_pytorch_tpc(name: str, tp_model: tp.TargetPlatformModel):
+    """
+    Generates a TargetPlatformCapabilities object with default operation sets to layers mapping.
+    Args:
+        name: Name of the TargetPlatformModel.
+        tp_model: TargetPlatformModel object.
+    Returns: a TargetPlatformCapabilities object for the given TargetPlatformModel.
+    """
+
+    pytorch_tpc = tp.TargetPlatformCapabilities(tp_model,
+                                                name=name,
+                                                version=TPC_VERSION)
+
+    with pytorch_tpc:
+        tp.OperationsSetToLayers("NoQuantization", [Dropout,
+                                                    Flatten,
+                                                    dropout,
+                                                    flatten,
+                                                    split,
+                                                    operator.getitem,
+                                                    reshape,
+                                                    unsqueeze,
+                                                    BatchNorm2d,
+                                                    chunk,
+                                                    unbind,
+                                                    torch.Tensor.size,
+                                                    permute,
+                                                    transpose,
+                                                    equal,
+                                                    gather,
+                                                    topk])
+
+        tp.OperationsSetToLayers("Conv", [Conv2d, ConvTranspose2d])
+        tp.OperationsSetToLayers("FullyConnected", [Linear])
+        tp.OperationsSetToLayers("AnyReLU", [torch.relu,
+                                             ReLU,
+                                             ReLU6,
+                                             LeakyReLU,
+                                             relu,
+                                             relu6,
+                                             leaky_relu,
+                                             tp.LayerFilterParams(Hardtanh, min_val=0),
+                                             tp.LayerFilterParams(hardtanh, min_val=0)])
+
+        tp.OperationsSetToLayers("Add", [operator.add, add])
+        tp.OperationsSetToLayers("Sub", [operator.sub, sub])
+        tp.OperationsSetToLayers("Mul", [operator.mul, mul])
+        tp.OperationsSetToLayers("Div", [operator.truediv, div])
+        tp.OperationsSetToLayers("PReLU", [PReLU, prelu])
+        tp.OperationsSetToLayers("Swish", [SiLU, silu, Hardswish, hardswish])
+        tp.OperationsSetToLayers("Sigmoid", [Sigmoid, sigmoid])
+        tp.OperationsSetToLayers("Tanh", [Tanh, tanh])
+
+    return pytorch_tpc

model_compression_toolkit/exporter/__init__.py

@@ -12,3 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
+from model_compression_toolkit.exporter.model_exporter.keras.keras_export_facade import keras_export_model, KerasExportMode
+from model_compression_toolkit.exporter.model_exporter.pytorch.pytorch_export_facade import PyTorchExportMode, pytorch_export_model
+from model_compression_toolkit.exporter.model_exporter.tflite.tflite_export_facade import tflite_export_model, TFLiteExportMode
+

model_compression_toolkit/exporter/model_exporter/__init__.py

@@ -13,15 +13,3 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common.constants import FOUND_TF, FOUND_TORCH
-
-if FOUND_TF:
-    from model_compression_toolkit.exporter.model_exporter.keras.keras_export_facade import \
-        keras_export_model, KerasExportMode
-    from model_compression_toolkit.exporter.model_exporter.tflite.tflite_export_facade import tflite_export_model, \
-        TFLiteExportMode
-
-if FOUND_TORCH:
-    from model_compression_toolkit.exporter.model_exporter.pytorch.pytorch_export_facade import PyTorchExportMode, \
-        pytorch_export_model
-

model_compression_toolkit/exporter/model_exporter/fw_agonstic/exporter.py

@@ -54,4 +54,4 @@ class Exporter:
         Convert model and export it to a given path.
 
         """
-        Logger.critical(f'Exporter {self.__class__} have to implement export method')
+        Logger.critical(f'Exporter {self.__class__} have to implement export method')  # pragma: no cover

model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py

@@ -18,23 +18,12 @@ import keras.models
 import keras.models
 import tensorflow as tf
 from keras.engine.base_layer import Layer
-from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_configs import \
-    NoOpQuantizeConfig
 
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.exporter.model_exporter.keras.base_keras_exporter import \
     BaseKerasExporter
-from model_compression_toolkit.exporter.model_wrapper.keras.builder.quantize_config_to_node import \
-    SUPPORTED_QUANTIZATION_CONFIG
-from model_compression_toolkit.exporter.model_wrapper.keras.extended_quantize_wrapper import ExtendedQuantizeWrapper
-from model_compression_toolkit.exporter.model_wrapper.keras.quantize_configs.activation_quantize_config import \
-    ActivationQuantizeConfig
-from model_compression_toolkit.exporter.model_wrapper.keras.quantize_configs.weights_activation_quantize_config \
-    import \
-    WeightsActivationQuantizeConfig
-from model_compression_toolkit.exporter.model_wrapper.keras.quantize_configs.weights_quantize_config import \
-    WeightsQuantizeConfig
-from model_compression_toolkit.exporter.model_wrapper.keras.quantizers.fq_quantizer import FakeQuantQuantizer
+from model_compression_toolkit.quantizers_infrastructure import KerasQuantizationWrapper
+
 
 
 class FakelyQuantKerasExporter(BaseKerasExporter):
@@ -83,7 +72,7 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
             assert self.is_layer_exportable_fn(layer), f'Layer {layer.name} is not exportable.'
 
             # If weights are quantized, use the quantized weight for the new built layer.
-            if type(layer.quantize_config) in [WeightsQuantizeConfig, WeightsActivationQuantizeConfig]:
+            if layer.is_weights_quantization:
                 new_layer = layer.layer.__class__.from_config(layer.layer.get_config())
                 with tf.name_scope(new_layer.name):
                     new_layer.build(layer.input_shape)
@@ -100,8 +89,10 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
                             # that should be quantized. First, extract 'kernel' from variable name, check if the
                             # quantize config contains this as an attribute for quantization. If so -
                             # Take the quantized weight from the quantize_config and set it to the new layer.
-                            if w.name.split('/')[-1].split(':')[0] in layer.quantize_config.get_config()['weight_attrs']:
-                                val = layer.quantize_config.get_weights_and_quantizers(layer.layer)[0][1].weight
+                            attribute_name = w.name.split('/')[-1].split(':')[0]
+                            if attribute_name in layer.weights_quantizers.keys():
+                                quantizer = layer.weights_quantizers.get(attribute_name)
+                                val = quantizer(qw)
                             else:
                                 val = qw
                     if val is None:
@@ -113,23 +104,16 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
 
                 # If activations are also quantized, wrap the layer back using ActivationQuantizeConfig
                 # from original wrapper (weights wrapping is no longer needed).
-                if isinstance(layer.quantize_config, WeightsActivationQuantizeConfig):
-                    new_layer = ExtendedQuantizeWrapper(new_layer, layer.quantize_config.act_config)
+                if layer.is_activation_quantization:
+                    new_layer = KerasQuantizationWrapper(layer=new_layer,
+                                                         activation_quantizers=layer.activation_quantizers)
 
                 return new_layer
 
             # If this is a layer with activation quantization only, just return it
             # as activation quantization in the fake-quant case uses the wrapper for quantization.
-            elif type(layer.quantize_config) in [ActivationQuantizeConfig]:
-                return layer
+            return layer
 
-            # Ideally we want in the case of no quantization to simply use the inner layer.
-            # But for some reason when using SNC we are having issues to use the inner layer.
-            # The clone_model method tries to reconstruct a model from the unwrapped configuration,
-            # but when we have two TFOpLambda (like in the case of SNC: add and pad) one after another,
-            # the output shape of the first one is in correct (it adds a new axis
-            elif isinstance(layer.quantize_config, NoOpQuantizeConfig):
-                return layer
 
         # clone each layer in the model and apply _unwrap_quantize_wrapper to layers wrapped with a QuantizeWrapper.
         self.exported_model = tf.keras.models.clone_model(self.model,
@@ -137,7 +121,7 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
                                                           clone_function=_unwrap_quantize_wrapper)
 
         if self.exported_model is None:
-            Logger.critical(f'Exporter can not save model as it is not exported')
+            Logger.critical(f'Exporter can not save model as it is not exported')  # pragma: no cover
 
         Logger.info(f'Exporting FQ Keras model to: {self.save_model_path}')
 
@@ -145,14 +129,3 @@ class FakelyQuantKerasExporter(BaseKerasExporter):
 
         return FakelyQuantKerasExporter.get_custom_objects()
 
-    @staticmethod
-    def get_custom_objects() -> Dict[str, type]:
-        """
-
-        Returns: A dictionary with objects for loading the exported model.
-
-        """
-        return {ExtendedQuantizeWrapper.__name__: ExtendedQuantizeWrapper,
-                ActivationQuantizeConfig.__name__: ActivationQuantizeConfig,
-                FakeQuantQuantizer.__name__: FakeQuantQuantizer,
-                NoOpQuantizeConfig.__name__: NoOpQuantizeConfig}

model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -15,44 +15,56 @@
 from enum import Enum
 from typing import Callable, Dict
 
-import keras
 from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import \
-    FakelyQuantKerasExporter
+from model_compression_toolkit.core.common.constants import FOUND_TF
 
 
 class KerasExportMode(Enum):
     FAKELY_QUANT = 0
 
 
-def keras_export_model(model: keras.models.Model,
-                       is_layer_exportable_fn: Callable,
-                       mode: KerasExportMode = KerasExportMode.FAKELY_QUANT,
-                       save_model_path: str = None) -> Dict[str, type]:
-    """
-    Prepare and return fully quantized model for export. Save exported model to
-    a path if passed.
+if FOUND_TF:
+    import keras
+    from model_compression_toolkit.exporter.model_wrapper.keras.validate_layer import is_keras_layer_exportable
+    from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import FakelyQuantKerasExporter
 
-    Args:
-        model: Model to export.
-        is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-        mode: Mode to export the model according to.
-        save_model_path: Path to save the model.
+    def keras_export_model(model: keras.models.Model,
+                           save_model_path: str,
+                           is_layer_exportable_fn: Callable = is_keras_layer_exportable,
+                           mode: KerasExportMode = KerasExportMode.FAKELY_QUANT) -> Dict[str, type]:
+        """
+        Export a Keras quantized model to h5 model.
+        The model will be saved to the path in save_model_path.
+        Mode can be used for different exported files. Currently, keras_export_model
+        supports KerasExportMode.FAKELY_QUANT (where weights and activations are
+        float fakely-quantized values).
 
-    Returns:
-        Custom objects dictionary needed to load the model.
+        Args:
+            model: Model to export.
+            is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
+            mode: Mode to export the model according to.
+            save_model_path: Path to save the model.
 
-    """
+        Returns:
+            Custom objects dictionary needed to load the model.
 
-    if mode == KerasExportMode.FAKELY_QUANT:
-        exporter = FakelyQuantKerasExporter(model,
-                                            is_layer_exportable_fn,
-                                            save_model_path)
+        """
 
-    else:
-        Logger.critical(
-            f'Unsupported mode was used {mode.name} to export Keras model. Please see API for supported modes.')
+        if mode == KerasExportMode.FAKELY_QUANT:
+            exporter = FakelyQuantKerasExporter(model,
+                                                is_layer_exportable_fn,
+                                                save_model_path)
 
-    exporter.export()
+        else:
+            Logger.critical(
+                f'Unsupported mode was used {mode.name} to '
+                f'export Keras model. Please see API for supported modes.')  # pragma: no cover
 
-    return exporter.get_custom_objects()
+        exporter.export()
+
+        return exporter.get_custom_objects()
+else:
+    def keras_export_model(*args, **kwargs):
+        Logger.error('Installing tensorflow and tensorflow_model_optimization is mandatory '
+                     'when using keras_export_model. '
+                     'Could not find some or all of TensorFlow packages.')  # pragma: no cover

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -19,7 +19,13 @@ import torch.nn
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
 from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
+from packaging import version
 
+# ONNX opset version 16 is supported from PyTorch 1.12
+if version.parse(torch.__version__) < version.parse("1.12"):
+    OPSET_VERSION = 15
+else:
+    OPSET_VERSION = 16
 
 class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
     """
@@ -57,7 +63,9 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
         Returns:
             Fake-quant PyTorch model.
         """
-        # assert self.is_layer_exportable_fn(layer), f'Layer {layer.name} is not exportable.'
+        for layer in self.model.children():
+            assert self.is_layer_exportable_fn(layer), f'Layer {layer.name} is not exportable.'
+
         model_input = to_torch_tensor(next(self.repr_dataset())[0])
 
         Logger.info(f"Exporting PyTorch fake quant onnx model: {self.save_model_path}")
@@ -65,7 +73,7 @@ class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
         torch.onnx.export(self.model,
                           model_input,
                           self.save_model_path,
-                          opset_version=13,
+                          opset_version=OPSET_VERSION,
                           verbose=False,
                           input_names=['input'],
                           output_names=['output'],

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_torchscript_pytorch_exporter.py

@@ -56,8 +56,12 @@ class FakelyQuantTorchScriptPyTorchExporter(BasePyTorchExporter):
         Returns:
             Fake-quant PyTorch model.
         """
-        # assert self.is_layer_exportable_fn(layer), f'Layer {layer.name} is not exportable.'
-        torch_traced = torch.jit.trace(self.model, to_torch_tensor(next(self.repr_dataset())))
+        for layer in self.model.children():
+            assert self.is_layer_exportable_fn(layer), f'Layer {layer} is not exportable.'
+
+        torch_traced = torch.jit.trace(self.model,
+                                       to_torch_tensor(next(self.repr_dataset())),
+                                       check_trace=True)
         self.exported_model = torch.jit.script(torch_traced)
         Logger.info(f"Exporting PyTorch torch script Model: {self.save_model_path}")
         torch.jit.save(self.exported_model, self.save_model_path)

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -15,13 +15,8 @@
 from enum import Enum
 from typing import Callable
 
-import torch.nn
-
 from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
-    FakelyQuantONNXPyTorchExporter
-from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import \
-    FakelyQuantTorchScriptPyTorchExporter
+from model_compression_toolkit.core.common.constants import FOUND_TORCH
 
 
 class PyTorchExportMode(Enum):
@@ -29,38 +24,56 @@ class PyTorchExportMode(Enum):
     FAKELY_QUANT_ONNX = 1
 
 
-def pytorch_export_model(model: torch.nn.Module,
-                         is_layer_exportable_fn: Callable,
-                         mode: PyTorchExportMode = PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT,
-                         save_model_path: str = None,
-                         repr_dataset: Callable = None) -> None:
-    """
-    Prepare and return fully quantized model for export. Save exported model to
-    a path if passed.
+if FOUND_TORCH:
+    import torch.nn
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import FakelyQuantONNXPyTorchExporter
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import FakelyQuantTorchScriptPyTorchExporter
+    from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
+
+    def pytorch_export_model(model: torch.nn.Module,
+                             save_model_path: str,
+                             repr_dataset: Callable,
+                             is_layer_exportable_fn: Callable = is_pytorch_layer_exportable,
+                             mode: PyTorchExportMode = PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT) -> None:
+        """
+        Export a PyTorch quantized model to a torchscript or onnx model.
+        The model will be saved to the path in save_model_path.
+        Mode can be used for different exported files. Currently, pytorch_export_model
+        supports PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT (where the exported model
+        is in a TorchScript format and its weights and activations are float fakely-quantized values),
+        and PyTorchExportMode.FakelyQuantONNX (where the exported model
+        is in an ONNX format and its weights and activations are float fakely-quantized values)
+
+        Args:
+            model: Model to export.
+            is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
+            mode: Mode to export the model according to.
+            save_model_path: Path to save the model.
+            repr_dataset: Representative dataset for tracing the pytorch model (mandatory for exporting it).
 
-    Args:
-        model: Model to export.
-        is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-        mode: Mode to export the model according to.
-        save_model_path: Path to save the model.
-        repr_dataset: Representative dataset for tracing the pytorch model (mandatory for exporting it).
+        """
 
-    """
+        if mode == PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT:
+            exporter = FakelyQuantTorchScriptPyTorchExporter(model,
+                                                             is_layer_exportable_fn,
+                                                             save_model_path,
+                                                             repr_dataset)
 
-    if mode == PyTorchExportMode.FAKELY_QUANT_TORCHSCRIPT:
-        exporter = FakelyQuantTorchScriptPyTorchExporter(model,
-                                                         is_layer_exportable_fn,
-                                                         save_model_path,
-                                                         repr_dataset)
+        elif mode == PyTorchExportMode.FAKELY_QUANT_ONNX:
+            exporter = FakelyQuantONNXPyTorchExporter(model,
+                                                      is_layer_exportable_fn,
+                                                      save_model_path,
+                                                      repr_dataset)
 
-    elif mode == PyTorchExportMode.FAKELY_QUANT_ONNX:
-        exporter = FakelyQuantONNXPyTorchExporter(model,
-                                                  is_layer_exportable_fn,
-                                                  save_model_path,
-                                                  repr_dataset)
+        else:
+            Logger.critical(
+                f'Unsupported mode was used {mode.name} to export PyTorch model. '
+                f'Please see API for supported modes.')  # pragma: no cover
 
-    else:
-        Logger.critical(
-            f'Unsupported mode was used {mode.name} to export PyTorch model. Please see API for supported modes.')
+        exporter.export()
 
-    exporter.export()
+else:
+    def pytorch_export_model(*args, **kwargs):
+        Logger.error('Installing torch is mandatory '
+                     'when using pytorch_export_model. '
+                     'Could not find PyTorch packages.')  # pragma: no cover

model_compression_toolkit/exporter/model_exporter/tflite/fakely_quant_tflite_exporter.py

@@ -18,8 +18,8 @@ from typing import Callable
 
 import keras.models
 import tensorflow as tf
-from keras.models import load_model
 
+from model_compression_toolkit.quantizers_infrastructure.inferable_infrastructure.keras.load_model import keras_load_quantized_model
 from model_compression_toolkit.core.common import Logger
 from model_compression_toolkit.exporter.model_exporter.keras.fakely_quant_keras_exporter import FakelyQuantKerasExporter
 
@@ -61,7 +61,8 @@ class FakelyQuantTFLiteExporter(FakelyQuantKerasExporter):
         custom_objects = FakelyQuantKerasExporter(self.model,
                                                   self.is_layer_exportable_fn,
                                                   tmp_h5_file).export()
-        model = load_model(tmp_h5_file, custom_objects)
+
+        model = keras_load_quantized_model(tmp_h5_file)
         os.remove(tmp_h5_file)
 
         self.exported_model = tf.lite.TFLiteConverter.from_keras_model(model).convert()