Quantization Code Move (#570)

* Moved:
- `modifier_quantization.py` to `sparseml.pytrorch.sparsification.quantization`
- Updated test location and usages
_ Style

* fix imports

* Style

* Moved `sparseml/pytorch/utils/quantization` to `sparseml/pytorch/sparsification/quantization/utils`

* Moved `sparseml/pytorch/utils/quantization` to `sparseml/pytorch/sparsification/quantization/utils`

* Style

* Style

* fix imports

* Style

* Review Comments

* fix quality + failing tests + cyclic dependency

* quality

Co-authored-by: Benjamin Fineran <bfineran@users.noreply.github.com>
Co-authored-by: Benjamin <ben@neuralmagic.com>

Author: 3 people

src/sparseml/pytorch/optim/__init__.py

@@ -34,7 +34,6 @@
 from .modifier_lr import *
 from .modifier_params import *
 from .modifier_pruning import *
-from .modifier_quantization import *
 from .modifier_regularizer import *
 from .optimizer import *
 from .sensitivity_as import *

src/sparseml/pytorch/sparsification/__init__.py

@@ -24,3 +24,4 @@
 from .info import *
 from .modifier_thinning import *
 from .pruning import *
+from .quantization.modifier_quantization import *

…l/pytorch/utils/quantization/__init__.py …/sparsification/quantization/__init__.pysrc/sparseml/pytorch/utils/quantization/__init__.py renamed to src/sparseml/pytorch/sparsification/quantization/__init__.py src/sparseml/pytorch/utils/quantization/__init__.py renamed to src/sparseml/pytorch/sparsification/quantization/__init__.py

@@ -1,3 +1,5 @@
+# flake8: noqa
+
 # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Tools for quantizing and exporting PyTorch models
-"""
-
-# flake8: noqa
-
 from .helpers import *
+from .modifier_quantization import *
 from .quantize_qat_export import *

…ml/pytorch/utils/quantization/helpers.py …h/sparsification/quantization/helpers.pysrc/sparseml/pytorch/utils/quantization/helpers.py renamed to src/sparseml/pytorch/sparsification/quantization/helpers.py src/sparseml/pytorch/utils/quantization/helpers.py renamed to src/sparseml/pytorch/sparsification/quantization/helpers.py

@@ -46,8 +46,7 @@
 
 from sparseml.optim import BaseModifier, ModifierProp
 from sparseml.pytorch.optim.modifier import PyTorchModifierYAML, ScheduledModifier
-from sparseml.pytorch.utils import BaseLogger, tensors_module_forward, tensors_to_device
-from sparseml.pytorch.utils.quantization import (
+from sparseml.pytorch.sparsification.quantization.helpers import (
     add_quant_dequant,
     configure_module_default_qconfigs,
     configure_module_qat_wrappers,
@@ -57,6 +56,7 @@
     prepare_embeddings_qat,
     remove_activation_qat_by_layer_name,
 )
+from sparseml.pytorch.utils import BaseLogger, tensors_module_forward, tensors_to_device
 from sparseml.sparsification import SparsificationTypes
 
 

…l/pytorch/optim/modifier_quantization.py …on/quantization/modifier_quantization.pysrc/sparseml/pytorch/optim/modifier_quantization.py renamed to src/sparseml/pytorch/sparsification/quantization/modifier_quantization.py src/sparseml/pytorch/optim/modifier_quantization.py renamed to src/sparseml/pytorch/sparsification/quantization/modifier_quantization.py

@@ -42,10 +42,6 @@
     script_model,
     trace_model,
 )
-from sparseml.pytorch.utils.quantization import (
-    quantize_torch_qat_export,
-    skip_onnx_input_quantize,
-)
 from sparseml.utils import clean_path, create_parent_dirs
 
 
@@ -493,10 +489,20 @@ def export_onnx(
 
     if convert_qat and is_quant_module:
         # overwrite exported model with fully quantized version
+        # import here to avoid cyclic dependency
+        from sparseml.pytorch.sparsification.quantization import (
+            quantize_torch_qat_export,
+        )
+
         quantize_torch_qat_export(model=file_path, output_file_path=file_path)
 
     if skip_input_quantize:
         try:
+            # import here to avoid cyclic dependency
+            from sparseml.pytorch.sparsification.quantization import (
+                skip_onnx_input_quantize,
+            )
+
             skip_onnx_input_quantize(file_path, file_path)
         except Exception as e:
             _LOGGER.warning(

…tils/quantization/quantize_qat_export.py …tion/quantization/quantize_qat_export.pysrc/sparseml/pytorch/utils/quantization/quantize_qat_export.py renamed to src/sparseml/pytorch/sparsification/quantization/quantize_qat_export.py src/sparseml/pytorch/utils/quantization/quantize_qat_export.py renamed to src/sparseml/pytorch/sparsification/quantization/quantize_qat_export.py

@@ -11,5 +11,3 @@
 # 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.
-
-# flake8: noqa

src/sparseml/pytorch/utils/exporter.py

@@ -18,7 +18,7 @@
 import torch
 
 from sparseml.pytorch.models import mobilenet, resnet50
-from sparseml.pytorch.utils.quantization import (
+from sparseml.pytorch.sparsification.quantization import (
     QATWrapper,
     add_quant_dequant,
     configure_module_default_qconfigs,

…l/pytorch/utils/quantization/__init__.py …/sparsification/quantization/__init__.pytests/sparseml/pytorch/utils/quantization/__init__.py renamed to tests/sparseml/pytorch/sparsification/quantization/__init__.py tests/sparseml/pytorch/utils/quantization/__init__.py renamed to tests/sparseml/pytorch/sparsification/quantization/__init__.py

@@ -17,7 +17,7 @@
 import pytest
 from torch.nn import Conv2d, Identity, Linear
 
-from sparseml.pytorch.optim import QuantizationModifier
+from sparseml.pytorch.sparsification import QuantizationModifier
 from tests.sparseml.pytorch.helpers import LinearNet, create_optim_sgd
 from tests.sparseml.pytorch.optim.test_modifier import ScheduledModifierTest