Add SGDOptimizer in the on-device training offline tooling (onnxblock) (

microsoft#17085) ### Description Adding SGDOptimizer to on device training onnxblock
JiazhouLvguan · Aug 18, 2023 · c0b6c6c · c0b6c6c
1 parent ee09a5f
commit c0b6c6c
Show file tree

Hide file tree

Showing 11 changed files with 266 additions and 123 deletions.
diff --git a/onnxruntime/test/testdata/training_api/adamw.onnx b/onnxruntime/test/testdata/training_api/adamw.onnx
diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc
@@ -1537,7 +1537,7 @@ void RegisterTrainingOpSchemas() {
              "This signal indicates if weight updates are skipped, applicable to gradient infinity check"
              " in mixed precision training. ",
              "T_BOOL", OpSchema::Optional)
-      .Output(0, "updated_flag", "Whether gradient is applied or not.", "T2")
+      .Output(0, "updated_flag", "Whether gradient is applied or not.", "T_BOOL")
       .Output(1, "updated_weights", "Sequence of weights after optimize.", "S_WEIGHT", OpSchema::Optional)
       .Output(2, "updated_momentums_1", "Sequence of momentum_1 after optimize.", "S_MOMENT", OpSchema::Optional)
       .Output(3, "updated_momentums_2", "Sequence of momentum_2 after optimize.", "S_MOMENT", OpSchema::Optional)

diff --git a/orttraining/orttraining/python/training/artifacts.py b/orttraining/orttraining/python/training/artifacts.py
@@ -33,6 +33,7 @@ class OptimType(Enum):
     """
 
     AdamW = 1
+    SGD = 2
 
 
 def generate_artifacts(
@@ -192,7 +193,8 @@ def _export_to_ort_format(model_path, output_dir, extra_options):
     logging.info("Optimizer enum provided: %s", optimizer.name)
 
     optim_model = None
-    optim_blocks = {OptimType.AdamW: onnxblock.optim.AdamW}
+    optim_blocks = {OptimType.AdamW: onnxblock.optim.AdamW, OptimType.SGD: onnxblock.optim.SGD}
+
     optim_block = optim_blocks[optimizer]()
     with onnxblock.empty_base():
         _ = optim_block(model_params)

diff --git a/orttraining/orttraining/python/training/onnxblock/optim/__init__.py b/orttraining/orttraining/python/training/onnxblock/optim/__init__.py
@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
-from onnxruntime.training.onnxblock.optim.optim import AdamW, ClipGradNorm
+from onnxruntime.training.onnxblock.optim.optim import SGD, AdamW, ClipGradNorm
 
-__all__ = ["AdamW", "ClipGradNorm"]
+__all__ = ["AdamW", "ClipGradNorm", "SGD"]