Miscellaneous updates to training artifact generation (#15315)

orttraining/orttraining/python/training/onnxblock/_graph_utils.py

@@ -72,3 +72,16 @@ def register_graph_outputs(model: onnx.ModelProto, output_names: Union[List[str]
     del model.graph.output[:]
 
     model.graph.output.extend(graph_outputs)
+
+
+def node_arg_exists(model: onnx.ModelProto, node_arg_name: str) -> bool:
+    """Returns True if the given node_arg_name exists in the model graph."""
+
+    for node in model.graph.node:
+        if node_arg_name in node.input:
+            return True
+
+        if node_arg_name in node.output:
+            return True
+
+    return False

orttraining/orttraining/python/training/onnxblock/_training_graph_utils.py

@@ -9,6 +9,29 @@ import onnx
 from onnxruntime.capi._pybind_state import GradientGraphBuilder, get_optimized_model
 
 
+def _disable_training_mode(model: onnx.ModelProto) -> None:
+    """Disables the training mode of the model by removing the training configuration."""
+
+    def disable_training_mode_dropout(node):
+        # Training mode is the third input of Dropout
+        if len(node.input) > 2:
+            node.input[2] = ""
+
+    def disable_training_mode_batchnorm(node):
+        # Training mode is an attribute of BatchNormalization
+        for attr in node.attribute:
+            if attr.name == "training_mode":
+                attr.i = 0
+
+    ops_to_disable_training_mode_func_map = {
+        "Dropout": disable_training_mode_dropout,
+        "BatchNormalization": disable_training_mode_batchnorm,
+    }
+    for node in model.graph.node:
+        if node.op_type in ops_to_disable_training_mode_func_map:
+            ops_to_disable_training_mode_func_map[node.op_type](node)
+
+
 def _reorder_outputs(model: onnx.ModelProto, user_output_names: List[str], requires_grad: Set[str]) -> None:
     """Reorders the outputs of the model to match the order of [user_outputs, gradients]"""
 
@@ -81,6 +104,7 @@ def build_gradient_graph(
 
     # At this point, eval model and training model diverge.
     eval_model = copy.deepcopy(model)
+    _disable_training_mode(eval_model)
 
     optimized_model = onnx.load_from_string(get_optimized_model(model.SerializeToString(), requires_grad))
 

orttraining/orttraining/python/training/onnxblock/blocks.py

@@ -5,7 +5,7 @@ import contextlib
 import copy
 import logging
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Any, List, Optional
 
 import onnx
 
@@ -361,3 +361,69 @@ class InputLike(Block):
         self.base.graph.input.append(cloned_input)
 
         return cloned_input.name
+
+
+class LabelEncoder(Block):
+    def __init__(
+        self,
+        default_float: float = 0.0,
+        default_int64: int = -1,
+        default_string: str = "_Unused",
+        keys_floats: Optional[List[float]] = None,
+        keys_int64s: Optional[List[int]] = None,
+        keys_strings: Optional[List[str]] = None,
+        values_floats: Optional[List[float]] = None,
+        values_int64s: Optional[List[int]] = None,
+        values_strings: Optional[List[str]] = None,
+    ):
+        super().__init__()
+
+        self._attributes = {
+            "default_float": default_float,
+            "default_int64": default_int64,
+            "default_string": default_string,
+        }
+
+        def _add_attributes(names: List[str], values: List[Any]):
+            for name, value in zip(names, values):
+                if value is not None:
+                    self._attributes[name] = value
+
+        _add_attributes(
+            ["keys_floats", "keys_int64s", "keys_strings", "values_floats", "values_int64s", "values_strings"],
+            [keys_floats, keys_int64s, keys_strings, values_floats, values_int64s, values_strings],
+        )
+
+    def build(self, label_encoder_input_name: str):
+        label_encoder_output_name = _graph_utils.generate_graph_name("label_encoder.output")
+        label_encoder_node = onnx.helper.make_node(
+            "LabelEncoder",
+            [label_encoder_input_name],
+            [label_encoder_output_name],
+            _graph_utils.generate_graph_name("LabelEncoder"),
+            domain="ai.onnx.ml",
+            **self._attributes,
+        )
+        self.base.graph.node.append(label_encoder_node)
+
+        return label_encoder_output_name
+
+
+class Cast(Block):
+    def __init__(self, to: onnx.TensorProto.DataType):
+        super().__init__()
+
+        self._to = to
+
+    def build(self, cast_input_name: str):
+        cast_output_name = _graph_utils.generate_graph_name("cast.output")
+        cast_node = onnx.helper.make_node(
+            "Cast",
+            [cast_input_name],
+            [cast_output_name],
+            _graph_utils.generate_graph_name("Cast"),
+            to=self._to,
+        )
+        self.base.graph.node.append(cast_node)
+
+        return cast_output_name

orttraining/orttraining/python/training/onnxblock/loss/loss.py

@@ -44,7 +44,10 @@ class MSELoss(blocks.Block):
             Returns a string of the output name from the loss
         """
 
-        return self._reduce(self._square(self._sub(loss_input_name, blocks.InputLike(loss_input_name)(target_name))))
+        if not _graph_utils.node_arg_exists(self.base, target_name):
+            target_name = blocks.InputLike(loss_input_name)(target_name)
+
+        return self._reduce(self._square(self._sub(loss_input_name, target_name)))
 
 
 class CrossEntropyLoss(blocks.Block):
@@ -84,15 +87,16 @@ class CrossEntropyLoss(blocks.Block):
         if self._weight is not None:
             self.base.graph.initializer.append(onnx.numpy_helper.from_array(self._weight, weight_name))
 
-        # create a new graph input. this is the labels input needed to compare
-        # the graph output against to calculate loss.
-        labels_input = copy.deepcopy(_graph_utils.get_output_from_output_name(self.base, scores_input_name))
-        labels_input.name = labels_name
-        labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT32
-        # if the predictions are (num_examples x num_classes)
-        # labels should be (num_examples x 1)
-        del labels_input.type.tensor_type.shape.dim[1]
-        self.base.graph.input.append(labels_input)
+        if not _graph_utils.node_arg_exists(self.base, labels_name):
+            # Create a new graph input. This is the labels input needed to compare
+            # the graph output against to calculate loss.
+            labels_input = copy.deepcopy(_graph_utils.get_output_from_output_name(self.base, scores_input_name))
+            labels_input.name = labels_name
+            labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT64
+            # If the predictions are (num_examples x num_classes)
+            # labels should be (num_examples,)
+            del labels_input.type.tensor_type.shape.dim[1]
+            self.base.graph.input.append(labels_input)
 
         loss_node_input_names = [scores_input_name, labels_name]
         if self._weight:
@@ -178,11 +182,12 @@ class BCEWithLogitsLoss(blocks.Block):
             onnx.helper.make_tensor(sub_ones_operand_name2, onnx.TensorProto.FLOAT, [1], [1.0])
         )
 
-        # Create a new graph input. This is the target input needed to compare
-        # the graph output against to calculate loss.
-        target_input = copy.deepcopy(_graph_utils.get_output_from_output_name(self.base, loss_input_name))
-        target_input.name = target_name
-        self.base.graph.input.append(target_input)
+        if not _graph_utils.node_arg_exists(self.base, target_name):
+            # Create a new graph input. This is the target input needed to compare
+            # the graph output against to calculate loss.
+            target_input = copy.deepcopy(_graph_utils.get_output_from_output_name(self.base, loss_input_name))
+            target_input.name = target_name
+            self.base.graph.input.append(target_input)
 
         sigmoid_output = self._sigmoid(loss_input_name)
         add_operand1 = self._mul(self._log(sigmoid_output), target_name)
@@ -235,4 +240,7 @@ class L1Loss(blocks.Block):
             Returns a string of the output name from the loss
         """
 
-        return self._reduce(self._abs(self._sub(loss_input_name, blocks.InputLike(loss_input_name)(target_name))))
+        if not _graph_utils.node_arg_exists(self.base, target_name):
+            target_name = blocks.InputLike(loss_input_name)(target_name)
+
+        return self._reduce(self._abs(self._sub(loss_input_name, target_name)))

orttraining/orttraining/test/python/orttraining_test_onnxblock.py

@@ -12,6 +12,7 @@ import torch
 import onnxruntime
 import onnxruntime.training.onnxblock as onnxblock
 from onnxruntime.capi import _pybind_state as C
+from onnxruntime.training import artifacts
 
 # PyTorch Module definitions
 
@@ -293,7 +294,7 @@ def test_crossentropy_loss_execution():
     ort_output_names = [onnx_model.graph.output[0].name]
     ort_inputs = {
         onnx_model.graph.input[0].name: _to_numpy(copy.deepcopy(x)),
-        onnx_model.graph.input[1].name: _to_numpy(copy.deepcopy(target).type(torch.int32)),
+        onnx_model.graph.input[1].name: _to_numpy(copy.deepcopy(target).type(torch.int64)),
     }
 
     def crossentropy_loss(prediction, target):
@@ -410,7 +411,7 @@ def test_crossentropy_loss_training_graph_execution():
     onnx_model, _ = simple_block.to_model_proto()
 
     ort_output_names = _get_training_ort_output_names(pt_model, onnx_model)
-    ort_inputs = _get_training_ort_inputs(x, target, pt_model, onnx_model, target_type=torch.int32)
+    ort_inputs = _get_training_ort_inputs(x, target, pt_model, onnx_model, target_type=torch.int64)
 
     def crossentropy_loss(prediction, target):
         loss = torch.nn.CrossEntropyLoss()
@@ -816,3 +817,101 @@ def test_grad_clipping_execution():
         # assert all the gradients are close
         for ort_grad, pt_param in zip(ort_outs[0], pt_model.parameters()):
             assert np.allclose(ort_grad, _to_numpy(pt_param.grad))
+
+
+def test_eval_model_has_no_training_mode_dropout():
+    class DropoutModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.dropout = torch.nn.Dropout(p=0.5)
+
+        def forward(self, x):
+            return self.dropout(x)
+
+    model = DropoutModel()
+    onnx_model = _get_onnx_model(model, (torch.randn(1, 3, 224, 224),))
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        artifacts.generate_artifacts(onnx_model, loss=artifacts.LossType.CrossEntropyLoss, artifact_directory=temp_dir)
+
+        eval_model = onnx.load(os.path.join(temp_dir, "eval_model.onnx"))
+
+        flag = False
+        for node in eval_model.graph.node:
+            if node.op_type == "Dropout":
+                assert not node.input[2]
+                flag = True
+
+        assert flag
+
+
+def test_eval_model_has_no_training_mode_batchnorm():
+    class BatchNormModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.batchnorm = torch.nn.BatchNorm2d(100)
+
+        def forward(self, x):
+            return self.batchnorm(x)
+
+    model = BatchNormModel()
+    onnx_model = _get_onnx_model(model, (torch.randn(20, 100, 35, 45),))
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        artifacts.generate_artifacts(onnx_model, loss=artifacts.LossType.CrossEntropyLoss, artifact_directory=temp_dir)
+
+        eval_model = onnx.load(os.path.join(temp_dir, "eval_model.onnx"))
+
+        flag = False
+        for node in eval_model.graph.node:
+            if node.op_type == "BatchNormalization":
+                for attr in node.attribute:
+                    if attr.name == "training_mode":
+                        assert attr.i == 0
+                        flag = True
+
+        assert flag
+
+
+def test_label_encoder_composition():
+    device = "cuda"
+    batch_size, input_size, hidden_size, output_size = 64, 784, 500, 10
+    _, base_model = _get_models(device, batch_size, input_size, hidden_size, output_size)
+    base_model.opset_import.append(
+        onnx.helper.make_opsetid("ai.onnx.ml", onnx.defs.onnx_opset_version()),
+    )
+
+    all_nodes = [node.op_type for node in base_model.graph.node]
+    assert "LabelEncoder" not in all_nodes
+
+    class SCELossWithLabelEncoder(onnxblock.ForwardBlock):
+        def __init__(self):
+            super().__init__()
+            self._loss = onnxblock.loss.CrossEntropyLoss()
+
+        def build(self, output_name):
+            keys_int64s = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+            values_int64s = [521, 522, 523, 524, 525, 526, 527, 528, 529, 530]
+
+            # Create a new graph input for the labels
+            labels_name = "labels"
+            labels_input = copy.deepcopy(self.base.graph.output[0])
+            labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT64
+            labels_input.name = labels_name
+            del labels_input.type.tensor_type.shape.dim[1]
+            self.base.graph.input.append(labels_input)
+
+            label_encoder = onnxblock.blocks.LabelEncoder(
+                default_int64=521, keys_int64s=keys_int64s, values_int64s=values_int64s
+            )
+
+            return self._loss(output_name, label_encoder(labels_name))
+
+    block = SCELossWithLabelEncoder()
+    model = None
+    with onnxblock.base(base_model):
+        _ = block(base_model.graph.output[0].name)
+        model = block.to_model_proto()
+
+    all_nodes = [node.op_type for node in model.graph.node]
+    assert "LabelEncoder" in all_nodes

orttraining/orttraining/training_api/module.cc

@@ -7,6 +7,7 @@
 #include "core/session/inference_session.h"
 #include "core/session/environment.h"
 #include "core/session/onnxruntime_session_options_config_keys.h"
+#include "core/graph/graph_utils.h"
 
 #include "orttraining/training_api/module.h"
 #include "orttraining/training_api/utils.h"
@@ -43,13 +44,17 @@ std::unordered_set<const Node*> GetReverseReachableNodes(Graph& inference_graph,
 }
 
 Status RemoveUnusedNodes(Graph& inference_graph, InlinedVector<const NodeArg*>& output_node_args) {
-  auto reachable_nodes = GetReverseReachableNodes(inference_graph, output_node_args);
+  const auto reachable_nodes = GetReverseReachableNodes(inference_graph, output_node_args);
 
   // Get all graph nodes and remove those that are not in the reachable nodes.
   GraphViewer graph_viewer(inference_graph);
-  for (auto& node : graph_viewer.Nodes()) {
-    if (!reachable_nodes.count(&node)) {
-      inference_graph.RemoveNode(node.Index());
+  const auto node_indices = graph_viewer.GetNodesInTopologicalOrder();
+  for (size_t idx = node_indices.size(); idx > 0; --idx) {
+    const NodeIndex node_index = idx - 1;
+    auto* node = inference_graph.GetNode(node_index);
+    if (!reachable_nodes.count(node)) {
+      graph_utils::RemoveNodeOutputEdges(inference_graph, *node);
+      inference_graph.RemoveNode(node_index);
     }
   }