update import

sony · Sep 3, 2024 · 882bff6 · 882bff6
1 parent bc729e9
commit 882bff6
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 8 deletions.
diff --git a/model_compression_toolkit/qat/keras/quantization_facade.py b/model_compression_toolkit/qat/keras/quantization_facade.py
@@ -24,7 +24,6 @@
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     MixedPrecisionQuantizationConfig
 from mct_quantizers import KerasActivationQuantizationHolder
-from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
 from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework import TargetPlatformCapabilities
 from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.ptq.runner import ptq_runner
@@ -34,6 +33,7 @@
     from tensorflow.keras.layers import Layer
     from tensorflow.keras.models import Model
 
+    from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation

diff --git a/model_compression_toolkit/trainable_infrastructure/__init__.py b/model_compression_toolkit/trainable_infrastructure/__init__.py
@@ -15,9 +15,10 @@
 
 from model_compression_toolkit.trainable_infrastructure.common.trainable_quantizer_config import TrainableQuantizerWeightsConfig, TrainableQuantizerActivationConfig
 from model_compression_toolkit.trainable_infrastructure.common.training_method import TrainingMethod
-from model_compression_toolkit.trainable_infrastructure.keras.base_keras_quantizer import BaseKerasTrainableQuantizer
-from model_compression_toolkit.trainable_infrastructure.keras.quantize_wrapper import KerasTrainableQuantizationWrapper
-from model_compression_toolkit.verify_packages import FOUND_TORCH
+from model_compression_toolkit.verify_packages import FOUND_TORCH, FOUND_TF
+if FOUND_TF:
+    from model_compression_toolkit.trainable_infrastructure.keras.base_keras_quantizer import BaseKerasTrainableQuantizer
+    from model_compression_toolkit.trainable_infrastructure.keras.quantize_wrapper import KerasTrainableQuantizationWrapper
 
 if FOUND_TORCH:
     from model_compression_toolkit.trainable_infrastructure.pytorch.base_pytorch_quantizer import BasePytorchTrainableQuantizer

diff --git a/model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py b/model_compression_toolkit/trainable_infrastructure/keras/base_keras_quantizer.py
@@ -81,8 +81,8 @@ def get_trainable_variables(self, group: VariableGroup) -> List[tf.Tensor]:
 
 
 else:
-    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):
+    class BaseKerasTrainableQuantizer(BaseTrainableQuantizer):     # pragma: no cover
         def __init__(self, *args, **kwargs):
             Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                             "BaseKerasTrainableQuantizer. The 'tensorflow' package is missing "
-                            "or is installed with a version higher than 2.15.")  # pragma: no cover
+                            "or is installed with a version higher than 2.15.")
diff --git a/model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py b/model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py
@@ -101,7 +101,7 @@ def convert_to_inferable_quantizers(self):
             return inferable_quantizers_wrapper
 
 else:
-    class KerasTrainableQuantizationWrapper:
+    class KerasTrainableQuantizationWrapper:    # pragma: no cover
         def __init__(self, *args, **kwargs):
             """
             Keras Quantization Wrapper takes a keras layer and quantizers and infer a quantized layer.
@@ -112,4 +112,4 @@ def __init__(self, *args, **kwargs):
             """
             Logger.critical("Tensorflow must be installed with a version of 2.15 or lower to use "
                             "KerasTrainableQuantizationWrapper. The 'tensorflow' package is missing "
-                            "or is installed with a version higher than 2.15.")  # pragma: no cover
+                            "or is installed with a version higher than 2.15.")