Add flag for push-button transition between v1 and v2 (#3479)

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
quic · Nov 13, 2024 · 9ad3254 · 9ad3254
1 parent 556e121
commit 9ad3254
Show file tree

Hide file tree

Showing 37 changed files with 807 additions and 697 deletions.
diff --git a/Docs/api_docs/torch_batchnorm_re_estimation.rst b/Docs/api_docs/torch_batchnorm_re_estimation.rst
@@ -24,7 +24,7 @@ Top-level APIs
 
 **API for BatchNorm fold to scale**
 
-.. autofunction:: aimet_torch.batch_norm_fold.fold_all_batch_norms_to_scale
+.. autofunction:: aimet_torch.v1.batch_norm_fold.fold_all_batch_norms_to_scale
 
 Code Example - BN-Reestimation
 ==============================

diff --git a/Docs/api_docs/torch_primitive_apis_cle.rst b/Docs/api_docs/torch_primitive_apis_cle.rst
@@ -36,7 +36,7 @@ Higher Level APIs for Cross Layer Equalization
 
 **API for Batch Norm Folding**
 
-.. autofunction:: aimet_torch.batch_norm_fold.fold_all_batch_norms
+.. autofunction:: aimet_torch.v1.batch_norm_fold.fold_all_batch_norms
 
 |
 
@@ -79,7 +79,7 @@ Lower Level APIs for Cross Layer Equalization
 
 **API for Batch Norm Folding**
 
-.. autofunction:: aimet_torch.batch_norm_fold.fold_given_batch_norms
+.. autofunction:: aimet_torch.v1.batch_norm_fold.fold_given_batch_norms
 
 |
 

diff --git a/Docs/torch_code_examples/bn_reestimation_example.py b/Docs/torch_code_examples/bn_reestimation_example.py
@@ -235,7 +235,7 @@ def call_bn_reestimation_apis(quant_sim, use_cuda):
 
     # Call fold_all_batch_norms_to_scale
 
-    from aimet_torch.batch_norm_fold import fold_all_batch_norms_to_scale
+    from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms_to_scale
 
     fold_all_batch_norms_to_scale(quant_sim)
 

diff --git a/Examples/torch/quantization/AMP.ipynb b/Examples/torch/quantization/AMP.ipynb
@@ -302,7 +302,7 @@
    },
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms\n",
     "\n",
     "_ = fold_all_batch_norms(model, input_shapes=(1, 3, 224, 224))"
    ]

diff --git a/Examples/torch/quantization/adaround.ipynb b/Examples/torch/quantization/adaround.ipynb
@@ -246,7 +246,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms\n",
     "\n",
     "_ = fold_all_batch_norms(model, input_shapes=(1, 3, 224, 224))"
    ]

diff --git a/Examples/torch/quantization/adaround.py b/Examples/torch/quantization/adaround.py
@@ -54,7 +54,7 @@
 import aimet_common
 from aimet_common.defs import QuantScheme
 from aimet_torch.v1.adaround.adaround_weight import Adaround, AdaroundParameters
-from aimet_torch.batch_norm_fold import fold_all_batch_norms
+from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms
 from aimet_torch.v1.quantsim import QuantizationSimModel
 
 # imports for data pipelines

diff --git a/Examples/torch/quantization/bn_reestimation.ipynb b/Examples/torch/quantization/bn_reestimation.ipynb
@@ -342,7 +342,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms_to_scale\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms_to_scale\n",
     "\n",
     "fold_all_batch_norms_to_scale(sim)"
    ]

diff --git a/Examples/torch/quantization/cle_bc.ipynb b/Examples/torch/quantization/cle_bc.ipynb
@@ -260,7 +260,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms\n",
     "\n",
     "_ = fold_all_batch_norms(model, input_shapes=(1, 3, 224, 224))"
    ]

diff --git a/Examples/torch/quantization/qat.ipynb b/Examples/torch/quantization/qat.ipynb
@@ -277,7 +277,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms\n",
     "\n",
     "_ = fold_all_batch_norms(model, input_shapes=(1, 3, 224, 224))"
    ]

diff --git a/Examples/torch/quantization/qat_range_learning.ipynb b/Examples/torch/quantization/qat_range_learning.ipynb
@@ -274,7 +274,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aimet_torch.batch_norm_fold import fold_all_batch_norms\n",
+    "from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms\n",
     "\n",
     "_ = fold_all_batch_norms(model, input_shapes=(1, 3, 224, 224))"
    ]

diff --git a/Examples/torch/quantization/range_learning.py b/Examples/torch/quantization/range_learning.py
@@ -54,7 +54,7 @@
 
 # imports for AIMET
 from aimet_common.defs import QuantScheme, EvalFunction
-from aimet_torch.batch_norm_fold import fold_all_batch_norms
+from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms
 from aimet_torch.cross_layer_equalization import equalize_model
 from aimet_torch.v1.quantsim import QuantizationSimModel
 

diff --git a/NightlyTests/torch/test_mixed_precision.py b/NightlyTests/torch/test_mixed_precision.py
@@ -54,7 +54,7 @@
 from aimet_common.defs import QuantizationDataType, CallbackFunc
 from aimet_torch.v1.quantsim import QuantizationSimModel
 from aimet_torch.model_preparer import prepare_model
-from aimet_torch.batch_norm_fold import fold_all_batch_norms
+from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms
 from aimet_torch.amp import utils as mixed_precision_utils
 from aimet_torch.amp.mixed_precision_algo import EvalCallbackFactory
 from aimet_torch.v1.qc_quantize_op import QcQuantizeWrapper

diff --git a/NightlyTests/torch/test_quantize_resnet18.py b/NightlyTests/torch/test_quantize_resnet18.py
@@ -57,7 +57,7 @@
 from models.supervised_classification_pipeline import create_stand_alone_supervised_classification_evaluator,\
     create_supervised_classification_trainer
 from aimet_torch.bn_reestimation import reestimate_bn_stats
-from aimet_torch.batch_norm_fold import fold_all_batch_norms_to_scale
+from aimet_torch.v1.batch_norm_fold import fold_all_batch_norms_to_scale
 from aimet_torch.model_preparer import prepare_model
 
 two_class_image_dir = './data/tiny-imagenet-2'

diff --git a/TrainingExtensions/torch/src/python/aimet_torch/adaround/__init__.py b/TrainingExtensions/torch/src/python/aimet_torch/adaround/__init__.py
@@ -34,5 +34,10 @@
 #
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
-""" Alias to legacy adaround """
-from ..v1.adaround import * # pylint: disable=wildcard-import, unused-wildcard-import
+""" Alias to v1/v2 adaround """
+from ..utils import _get_default_api
+
+if _get_default_api() == "v1":
+    from ..v1.adaround import * # pylint: disable=wildcard-import, unused-wildcard-import
+else:
+    from ..v2.adaround import * # pylint: disable=wildcard-import, unused-wildcard-import
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/adaround/adaround_weight.py b/TrainingExtensions/torch/src/python/aimet_torch/adaround/adaround_weight.py
@@ -34,13 +34,18 @@
 #
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
-""" Alias to legacy adaround.adaround_weight """
-from ..v1.adaround.adaround_weight import * # pylint: disable=wildcard-import, unused-wildcard-import
+""" Alias to v1/v2 adaround.adaround_weight """
+from ..utils import _get_default_api
 
-from ..utils import _warn_replaced_in_v2
-from ..v1.adaround import adaround_weight as _v1_api
-from ..v2 import adaround as _v2_api
+if _get_default_api() == "v1":
+    from ..v1.adaround.adaround_weight import * # pylint: disable=wildcard-import, unused-wildcard-import
 
-_warn_replaced_in_v2(__name__,
-                     v2_new_api=_v2_api.__name__,
-                     v1_legacy_api=_v1_api.__name__)
+    from ..utils import _warn_replaced_in_v2
+    from ..v1.adaround import adaround_weight as _v1_api
+    from ..v2 import adaround as _v2_api
+
+    _warn_replaced_in_v2(__name__,
+                         v2_new_api=_v2_api.__name__,
+                         v1_legacy_api=_v1_api.__name__)
+else:
+    from ..v2.adaround import * # pylint: disable=wildcard-import, unused-wildcard-import
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/arch_checker/arch_checker_rules.py b/TrainingExtensions/torch/src/python/aimet_torch/arch_checker/arch_checker_rules.py
@@ -50,7 +50,7 @@
 from aimet_common.utils import AimetLogger
 
 from aimet_torch.meta.connectedgraph import ConnectedGraph
-from aimet_torch.batch_norm_fold import find_standalone_batchnorm_ops
+from aimet_torch.v1.batch_norm_fold import find_standalone_batchnorm_ops
 
 logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Utils)
 

diff --git a/TrainingExtensions/torch/src/python/aimet_torch/auto_quant.py b/TrainingExtensions/torch/src/python/aimet_torch/auto_quant.py
@@ -34,13 +34,18 @@
 #
 #  @@-COPYRIGHT-END-@@
 # =============================================================================
-""" Alias to legacy auto_quant """
-from .v1.auto_quant import * # pylint: disable=wildcard-import, unused-wildcard-import
+""" Alias to v1/v2 auto_quant """
+from .utils import _get_default_api
 
-from .utils import _warn_replaced_in_v2
-from .v1 import auto_quant as _v1_api
-from .v2 import auto_quant as _v2_api
+if _get_default_api() == "v1":
+    from .v1.auto_quant import * # pylint: disable=wildcard-import, unused-wildcard-import
 
-_warn_replaced_in_v2(__name__,
-                     v2_new_api=_v2_api.__name__,
-                     v1_legacy_api=_v1_api.__name__)
+    from .utils import _warn_replaced_in_v2
+    from .v1 import auto_quant as _v1_api
+    from .v2 import auto_quant as _v2_api
+
+    _warn_replaced_in_v2(__name__,
+                         v2_new_api=_v2_api.__name__,
+                         v1_legacy_api=_v1_api.__name__)
+else:
+    from .v2.auto_quant import * # pylint: disable=wildcard-import, unused-wildcard-import