Move onnxruntime.experiment to onnxruntime.training namespace (#5045)

cmake/onnxruntime_python.cmake

@@ -171,16 +171,16 @@ endif()
 
 if (onnxruntime_ENABLE_TRAINING)
   file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS
-    "${ORTTRAINING_SOURCE_DIR}/python/training/*.py"
+    "${ORTTRAINING_SOURCE_DIR}/python/deprecated/*.py"
   )
   file(GLOB onnxruntime_python_root_srcs CONFIGURE_DEPENDS
-    "${ORTTRAINING_SOURCE_DIR}/python/experimental/*.py"
+    "${ORTTRAINING_SOURCE_DIR}/python/training/*.py"
   )
   file(GLOB onnxruntime_python_amp_srcs CONFIGURE_DEPENDS
-    "${ORTTRAINING_SOURCE_DIR}/python/experimental/amp/*.py"
+    "${ORTTRAINING_SOURCE_DIR}/python/training/amp/*.py"
   )
   file(GLOB onnxruntime_python_optim_srcs CONFIGURE_DEPENDS
-    "${ORTTRAINING_SOURCE_DIR}/python/experimental/optim/*.py"
+    "${ORTTRAINING_SOURCE_DIR}/python/training/optim/*.py"
   )
   file(GLOB onnxruntime_python_train_tools_srcs CONFIGURE_DEPENDS
     "${REPO_ROOT}/tools/python/register_custom_ops_pytorch_exporter.py"
@@ -280,21 +280,24 @@ add_custom_command(
 if (onnxruntime_ENABLE_TRAINING)
   add_custom_command(
     TARGET onnxruntime_pybind11_state POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental
-    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/amp
-    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/optim
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/amp
+    COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/optim
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_capi_training_srcs}
+        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/capi/training/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_root_srcs}
-        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/
+        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_amp_srcs}
-        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/amp/
+        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/amp/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_optim_srcs}
-        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/optim/
+        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/optim/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_train_tools_srcs}
-        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/experimental/
+        $<TARGET_FILE_DIR:${test_data_target}>/onnxruntime/training/
   )
 endif()
 

orttraining/orttraining/python/deprecated/__init__.py

@@ -4,7 +4,3 @@
 # --------------------------------------------------------------------------
 from onnxruntime.capi._pybind_state import TrainingParameters
 from onnxruntime.capi.training.training_session import TrainingSession
-
-from .orttrainer_options import ORTTrainerOptions
-from .orttrainer import ORTTrainer, TrainStepInfo
-from . import amp, optim, model_desc_validation

orttraining/orttraining/python/ort_trainer.py

@@ -9,7 +9,7 @@ import torch
 import torch.nn
 import torch.onnx
 import onnxruntime as ort
-from ..experimental import postprocess
+from ..training import postprocess
 from distutils.version import LooseVersion
 import warnings
 
@@ -347,7 +347,7 @@ def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, op
     sample_inputs_copy = copy.deepcopy(sample_inputs)
 
     # Enable contrib ops export from PyTorch
-    from onnxruntime.experimental import register_custom_ops_pytorch_exporter
+    from onnxruntime.training import register_custom_ops_pytorch_exporter
     register_custom_ops_pytorch_exporter.register_custom_op()
 
     torch.onnx._export(model, tuple(sample_inputs_copy), f,

orttraining/orttraining/python/training/__init__.py

@@ -4,3 +4,7 @@
 # --------------------------------------------------------------------------
 from onnxruntime.capi._pybind_state import TrainingParameters
 from onnxruntime.capi.training.training_session import TrainingSession
+
+from .orttrainer_options import ORTTrainerOptions
+from .orttrainer import ORTTrainer, TrainStepInfo
+from . import amp, checkpoint, optim, model_desc_validation

orttraining/orttraining/python/training/orttrainer.py

@@ -482,15 +482,16 @@ class ORTTrainer(object):
         # Deepcopy inputs, since input values may change after model run.
         sample_inputs_copy = copy.deepcopy(sample_inputs)
 
-        from onnxruntime.experimental import register_custom_ops_pytorch_exporter
+        # Handle contrib OPs support
+        from onnxruntime.training import register_custom_ops_pytorch_exporter
         if self.options._internal_use.enable_onnx_contrib_ops:
             # Enable contrib ops export from PyTorch
             register_custom_ops_pytorch_exporter.register_custom_op()
         else:
-            # unregister contrib ops, if they were registered in previous calls
+            # Unregister contrib ops, if they were registered in previous calls
             register_custom_ops_pytorch_exporter.unregister_custom_op()
 
-
+        # Export torch.nn.Module to ONNX
         torch.onnx._export(model, tuple(sample_inputs_copy), f,
                            input_names=[input.name for input in self.model_desc.inputs],
                            output_names=[output.name for output in self.model_desc.outputs],

orttraining/orttraining/test/python/_test_helpers.py

@@ -4,7 +4,7 @@ import sys
 import torch
 
 from numpy.testing import assert_allclose
-from onnxruntime.experimental import orttrainer
+from onnxruntime.training import orttrainer
 from onnxruntime.capi.ort_trainer import ORTTrainer as Legacy_ORTTrainer
 
 

orttraining/orttraining/test/python/orttraining_test_orttrainer_bert_toy_onnx.py

@@ -13,7 +13,7 @@ from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\
                                          ModelDescription as Legacy_ModelDescription,\
                                          LossScaler as Legacy_LossScaler,\
                                          ORTTrainer as Legacy_ORTTrainer
-from onnxruntime.experimental import _utils, amp, checkpoint, optim, orttrainer, TrainStepInfo,\
+from onnxruntime.training import _utils, amp, checkpoint, optim, orttrainer, TrainStepInfo,\
                                       model_desc_validation as md_val,\
                                       orttrainer_options as orttrainer_options
 

orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py

@@ -13,7 +13,7 @@ from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\
                                          ModelDescription as Legacy_ModelDescription,\
                                          LossScaler as Legacy_LossScaler,\
                                          ORTTrainer as Legacy_ORTTrainer
-from onnxruntime.experimental import _utils, amp, checkpoint, optim, orttrainer, TrainStepInfo,\
+from onnxruntime.training import _utils, amp, checkpoint, optim, orttrainer, TrainStepInfo,\
                                       model_desc_validation as md_val,\
                                       orttrainer_options as orttrainer_options
 import _test_commons,_test_helpers

orttraining/orttraining/test/python/orttraining_test_utils.py

@@ -5,10 +5,10 @@ from onnxruntime.capi.ort_trainer import ORTTrainer, IODescription
 from orttraining_test_data_loader import create_ort_test_dataloader, BatchArgsOption, split_batch
 from orttraining_test_bert_postprocess import postprocess_model
 
-from onnxruntime.experimental import _utils, amp, optim, orttrainer, TrainStepInfo,\
+from onnxruntime.training import _utils, amp, optim, orttrainer, TrainStepInfo,\
                                       model_desc_validation as md_val,\
                                       orttrainer_options as orttrainer_options
-from onnxruntime.experimental.optim import _LRScheduler
+from onnxruntime.training.optim import _LRScheduler
 
 def warmup_cosine(x, warmup=0.002):
     if x < warmup:
@@ -177,4 +177,3 @@ def run_test(model, model_desc, device, args, gradient_accumulation_steps, fp16,
         outputs = model.eval_step(*args, **kwargs)
 
     return (output.cpu().numpy() for output in outputs)
-

orttraining/orttraining/test/python/orttraining_transformer_trainer.py

@@ -24,10 +24,10 @@ import onnxruntime
 from orttraining_test_bert_postprocess import postprocess_model
 from onnxruntime.capi.ort_trainer import ORTTrainer, LossScaler, ModelDescription, IODescription
 
-from onnxruntime.experimental import _utils, amp, optim, orttrainer, TrainStepInfo,\
+from onnxruntime.training import _utils, amp, optim, orttrainer, TrainStepInfo,\
                                       model_desc_validation as md_val,\
                                       orttrainer_options as orttrainer_options
-from onnxruntime.experimental.optim import LinearWarmupLRScheduler, _LRScheduler
+from onnxruntime.training.optim import LinearWarmupLRScheduler, _LRScheduler
 
 try:
     from torch.utils.tensorboard import SummaryWriter

setup.py

@@ -238,11 +238,10 @@ packages = [
     'onnxruntime.quantization.operators',
 ]
 
-# TODO: thiagofc: Temporary 'experimental' namespace for new PyTorch front-end
 if '--enable_training' in sys.argv:
-    packages.extend(['onnxruntime.experimental',
-                     'onnxruntime.experimental.amp',
-                     'onnxruntime.experimental.optim'])
+    packages.extend(['onnxruntime.training',
+                     'onnxruntime.training.amp',
+                     'onnxruntime.training.optim'])
     sys.argv.remove('--enable_training')
 
 package_data = {}