From 6b755debbc440e4bf20af0ed18348d1a7a9b94c2 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Tue, 4 Apr 2023 20:09:51 -0700 Subject: [PATCH] Miscellaneous updates to training artifact generation (#15315) --- .../python/training/onnxblock/_graph_utils.py | 13 +++ .../onnxblock/_training_graph_utils.py | 24 ++++ .../python/training/onnxblock/blocks.py | 68 +++++++++++- .../python/training/onnxblock/loss/loss.py | 40 ++++--- .../test/python/orttraining_test_onnxblock.py | 103 +++++++++++++++++- .../orttraining/training_api/module.cc | 13 ++- 6 files changed, 238 insertions(+), 23 deletions(-) diff --git a/orttraining/orttraining/python/training/onnxblock/_graph_utils.py b/orttraining/orttraining/python/training/onnxblock/_graph_utils.py index b69c96df55..b6e223fa8d 100644 --- a/orttraining/orttraining/python/training/onnxblock/_graph_utils.py +++ b/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 diff --git a/orttraining/orttraining/python/training/onnxblock/_training_graph_utils.py b/orttraining/orttraining/python/training/onnxblock/_training_graph_utils.py index 6210a295f6..796267ea3a 100644 --- a/orttraining/orttraining/python/training/onnxblock/_training_graph_utils.py +++ b/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)) diff --git a/orttraining/orttraining/python/training/onnxblock/blocks.py b/orttraining/orttraining/python/training/onnxblock/blocks.py index 4ef9083243..f4b62a5e9d 100644 --- a/orttraining/orttraining/python/training/onnxblock/blocks.py +++ b/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 diff --git a/orttraining/orttraining/python/training/onnxblock/loss/loss.py b/orttraining/orttraining/python/training/onnxblock/loss/loss.py index 6bf73c6aac..2ca848fa3f 100644 --- a/orttraining/orttraining/python/training/onnxblock/loss/loss.py +++ b/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))) diff --git a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py b/orttraining/orttraining/test/python/orttraining_test_onnxblock.py index 4b8ea0f673..ce919f4706 100644 --- a/orttraining/orttraining/test/python/orttraining_test_onnxblock.py +++ b/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 diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc index 042b729016..67185b488e 100644 --- a/orttraining/orttraining/training_api/module.cc +++ b/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 GetReverseReachableNodes(Graph& inference_graph, } Status RemoveUnusedNodes(Graph& inference_graph, InlinedVector& 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); } }