From 92a8c650ada604d3b485f33c6b2c8323f89b68a3 Mon Sep 17 00:00:00 2001 From: Rayan-Krishnan Date: Fri, 11 Sep 2020 12:16:07 -0700 Subject: [PATCH] [Debuggability] Add feature to ORTTrainer Frontend (#5124) * add option, feature to orttrainer and test * address comments * minor fixes * further address comments * minor changes Co-authored-by: Rayan Krishnan --- .../orttraining/python/training/orttrainer.py | 40 ++++++++++++++ .../python/training/orttrainer_options.py | 11 ++++ .../python/orttraining_test_debuggability.py | 52 +++++++++++++++++++ .../orttraining_test_orttrainer_frontend.py | 3 +- 4 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 orttraining/orttraining/test/python/orttraining_test_debuggability.py diff --git a/orttraining/orttraining/python/training/orttrainer.py b/orttraining/orttraining/python/training/orttrainer.py index ba8a89327f..76b8b6d159 100644 --- a/orttraining/orttraining/python/training/orttrainer.py +++ b/orttraining/orttraining/python/training/orttrainer.py @@ -285,6 +285,40 @@ class ORTTrainer(object): with open(path, "wb") as f: f.write(self._onnx_model.SerializeToString()) + def _debug_model_export(self, input): + from onnx import helper, TensorProto, numpy_helper + import numpy as np + from numpy.testing import assert_allclose + import _test_helpers + onnx_model_copy = copy.deepcopy(self._onnx_model) + + # Mute the dropout nodes + dropout_nodes = [n for n in onnx_model_copy.graph.node if n.op_type == 'Dropout'] + for node in dropout_nodes: + ratio_node = [n for n in onnx_model_copy.graph.node if node.input[1] in n.output][0] + training_mode_node = [n for n in onnx_model_copy.graph.node if node.input[2] in n.output][0] + + training_mode_node.attribute.pop() + ratio_node.attribute.pop() + new_training_mode_arr = np.array(False, dtype=bool) + new_ratio_arr = np.array(0.0, dtype=np.float32) + new_training_mode = numpy_helper.from_array(new_training_mode_arr) + new_ratio = numpy_helper.from_array(new_ratio_arr) + training_mode_node.attribute.add().t.CopyFrom(new_training_mode) + ratio_node.attribute.add().t.CopyFrom(new_ratio) + training_mode_node.attribute[0].type = 4 + ratio_node.attribute[0].type = 4 + training_mode_node.attribute[0].name = "value" + ratio_node.attribute[0].name = "value" + + _inference_sess = ort.InferenceSession(onnx_model_copy.SerializeToString()) + inf_inputs = {} + for i, input_elem in enumerate(input): + inf_inputs[_inference_sess.get_inputs()[i].name] = input_elem.cpu().numpy() + _inference_outs = _inference_sess.run(None, inf_inputs) + for torch_item, ort_item in zip(self.torch_sample_outputs, _inference_outs): + assert_allclose(torch_item, ort_item, rtol=1e-2, atol=1e-6) + def train_step(self, *args, **kwargs): r"""Train step method @@ -303,6 +337,11 @@ class ORTTrainer(object): if self._onnx_model is None: sample_input = self._prepare_model_input(self.model_desc.inputs, None, None, *args, **kwargs) self._init_onnx_model(sample_input) + + # Debug Model Export if indicated + if self.options.debug.check_model_export: + self._debug_model_export(sample_input) + # Prepare inputs+lr and output descriptions inputs_desc = self._model_desc_inputs_with_lr @@ -461,6 +500,7 @@ class ORTTrainer(object): warnings.warn("This model cannot be deep copied (or pickled), which is a required step for stateful models to be properly exported to ONNX." " Compute will continue, but unexpected results may occur!") sample_outputs = model_copy(*sample_inputs_copy) + self.torch_sample_outputs = sample_outputs model.train() if isinstance(sample_outputs, torch.Tensor): diff --git a/orttraining/orttraining/python/training/orttrainer_options.py b/orttraining/orttraining/python/training/orttrainer_options.py index 3f9bdb768b..511761c9e8 100644 --- a/orttraining/orttraining/python/training/orttrainer_options.py +++ b/orttraining/orttraining/python/training/orttrainer_options.py @@ -144,6 +144,10 @@ class ORTTrainerOptions(object): 'type' : 'boolean', 'default' : False }, + 'check_model_export' : { + 'type' : 'boolean', + 'default' : False + }, } }, '_internal_use' : { @@ -225,6 +229,9 @@ class ORTTrainerOptions(object): debug options debug.deterministic_compute (bool, default is False) forces compute to be deterministic accross runs + debug.check_model_export (bool, default is False) + compares PyTorch outputs with InferenceSession outputs before the first + train step to ensure successful model export _internal_use (dict): internal options, possibly undocumented, that might be removed without notice _internal_use.enable_internal_postprocess (bool, default is True): @@ -450,6 +457,10 @@ _ORTTRAINER_OPTIONS_SCHEMA = { 'type': 'boolean', 'default': False }, + 'check_model_export': { + 'type': 'boolean', + 'default': False + }, } }, '_internal_use': { diff --git a/orttraining/orttraining/test/python/orttraining_test_debuggability.py b/orttraining/orttraining/test/python/orttraining_test_debuggability.py new file mode 100644 index 0000000000..4353174c38 --- /dev/null +++ b/orttraining/orttraining/test/python/orttraining_test_debuggability.py @@ -0,0 +1,52 @@ +import inspect +import onnx +import os +import pytest +import torch +import torchvision + +from numpy.testing import assert_allclose + +from onnxruntime import set_seed +from onnxruntime.capi.ort_trainer import IODescription as Legacy_IODescription,\ + ModelDescription as Legacy_ModelDescription,\ + LossScaler as Legacy_LossScaler,\ + ORTTrainer as Legacy_ORTTrainer +from onnxruntime.training import _utils, amp, optim, orttrainer, TrainStepInfo,\ + model_desc_validation as md_val,\ + orttrainer_options as orttrainer_options + +from orttraining_test_orttrainer_frontend import _load_pytorch_transformer_model + +import _test_helpers + + +############################################################################### +# Testing starts here ######################################################### +############################################################################### + + +@pytest.mark.parametrize("seed, device", [ + (24, 'cuda'), +]) +def testORTTransformerModelExport(seed, device): + # Common setup + optim_config = optim.LambConfig() + opts = orttrainer.ORTTrainerOptions({ + 'debug' : { + 'check_model_export': True, + }, + 'device' : { + 'id' : device, + } + }) + + # Setup for the first ORTTRainer run + torch.manual_seed(seed) + set_seed(seed) + model, model_desc, my_loss, batcher_fn, train_data, val_data, _ = _load_pytorch_transformer_model(device) + first_trainer = orttrainer.ORTTrainer(model, model_desc, optim_config, loss_fn=my_loss, options=opts) + data, targets = batcher_fn(train_data, 0) + _ = first_trainer.train_step(data, targets) + assert first_trainer._onnx_model is not None + diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py index 6071fbb742..fcef3af1bd 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py +++ b/orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py @@ -100,7 +100,8 @@ def testORTTrainerOptionsDefaultValues(test_input): 'invertible_layer_norm_gradient': False, }, 'debug': { - 'deterministic_compute': False + 'deterministic_compute': False, + 'check_model_export': False }, '_internal_use': { 'enable_internal_postprocess': True,