Adds classes to train an ONNX gradient with TrainingAgent

_doc/bench/bench_orttraining_nn_gpu.py

@@ -35,8 +35,9 @@
 from sklearn.neural_network import MLPRegressor
 from sklearn.metrics import mean_squared_error
 from mlprodict.onnx_conv import to_onnx
-from onnxcustom.training import add_loss_output, get_train_initializer
-from onnxcustom.training.optimizers import OrtGradientOptimizer
+from onnxcustom.training import (
+    add_loss_output, get_train_initializer,
+    OrtGradientOptimizer)
 
 
 def benchmark(N=1000, n_features=20, hidden_layer_sizes="25,25", max_iter=1000,

_doc/examples/plot_orttraining_benchmark.py

@@ -28,8 +28,9 @@
 from sklearn.model_selection import train_test_split
 from sklearn.neural_network import MLPRegressor
 from mlprodict.onnx_conv import to_onnx
-from onnxcustom.training import add_loss_output, get_train_initializer
-from onnxcustom.training.optimizers import OrtGradientOptimizer
+from onnxcustom.training import (
+    add_loss_output, get_train_initializer,
+    OrtGradientOptimizer)
 
 
 X, y = make_regression(2000, n_features=100, bias=2)

_unittests/ut_training/test_data_loader.py

@@ -7,7 +7,7 @@
 from pyquickhelper.pycode import ExtTestCase
 from sklearn.datasets import make_regression
 from onnxruntime import OrtValue
-from onnxcustom.training.data_loader import OrtDataLoader
+from onnxcustom.training import OrtDataLoader
 
 
 class TestDataLoadeer(ExtTestCase):

_unittests/ut_training/test_optimizers.py

@@ -22,8 +22,8 @@ class TestOptimizers(ExtTestCase):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -78,8 +78,8 @@ def test_ort_gradient_optimizers_use_ort(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_optimal_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -134,8 +134,8 @@ def test_ort_gradient_optimizers_optimal_use_ort(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_evaluation_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -164,8 +164,8 @@ def test_ort_gradient_optimizers_evaluation_use_numpy(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_evaluation_use_ort(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)

onnxcustom/training/__init__.py

@@ -1,6 +1,8 @@
 """
 @file
-@brief Shortcuts to *training*.
+@brief Shortcuts to *orttraining*.
 """
 
+from .data_loader import OrtDataLoader  # noqa
+from .optimizers import OrtGradientOptimizer  # noqa
 from .orttraining import add_loss_output, get_train_initializer  # noqa

onnxcustom/training/optimizers.py

@@ -1,6 +1,7 @@
 """
 @file
-@brief Helper for :epkg:`onnxruntime-training`.
+@brief Train a machine learned model
+with :epkg:`onnxruntime-training`.
 """
 import inspect
 import numpy

onnxcustom/training/ortgradient.py

@@ -0,0 +1,4 @@
+"""
+@file
+@brief Helpers for :epkg:`onnxruntime-training`.
+"""

onnxcustom/training/orttraining.py

@@ -1,6 +1,6 @@
 """
 @file
-@brief Helper for :epkg:`onnxruntime-training`.
+@brief Manipulate ONNX graph to train a model.
 """
 from onnx.helper import (
     make_node, make_graph, make_model, make_tensor_value_info,