From e71e08851a84c4c6417c04f79710ebb6e6b7bee6 Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Thu, 8 Oct 2020 18:07:32 -0700 Subject: [PATCH] Basic plumbing for backward pass. Not fully working --- .../orttraining/python/training/ortmodule.py | 118 +++++++++-------- .../orttraining_test_ortmodule_basic.py | 49 +++++++- ...ng_test_ortmodule_basic_transform_model.py | 116 ++++++++--------- .../mnist/{mnist_training.py => ort_mnist.py} | 46 ++++--- samples/python/mnist/pytorch_mnist.py | 119 ++++++++++++++++++ 5 files changed, 316 insertions(+), 132 deletions(-) rename samples/python/mnist/{mnist_training.py => ort_mnist.py} (78%) create mode 100644 samples/python/mnist/pytorch_mnist.py diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 4695a6afc8..bee6d7e7be 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -17,55 +17,37 @@ class ORTModule(torch.nn.Module): def __init__(self, module): print(f'ORTModule.__init__() was called') super(ORTModule, self).__init__() - # User will interact with it (debugging, etc) self._original_module = module # Forward pass self._onnx_forward = None + self._forward_session = None self._onnx_forward_initializers_desc = [] self._onnx_forward_inputs_desc = [] self._onnx_forward_outputs_desc = [] # Backward pass self._onnx_backward = None + self._backward_session = None def forward(self, *input, **kwargs): print(f'ORTModule.forward() was called') if not self._onnx_forward: original_forward_graph = ORTModule._get_forward_graph(self._original_module, *input, **kwargs) - # gradient_graph = ORTModule._build_gradient_graph(original_forward_graph) + gradient_graph = ORTModule._build_gradient_graph(original_forward_graph) + # TODO: Remove manual split after MVP # self.forward_graph, self.backward_graph = ORTModule._split_forward_and_backward(gradient_graph) - self._onnx_forward = original_forward_graph # TODO: hard-coding for MVP - # import pdb; pdb.set_trace() - self.forward_session = onnxruntime.InferenceSession(self._onnx_forward.SerializeToString()) + self._onnx_forward = original_forward_graph # TODO: hard-coding for MVP + self._onnx_backward = gradient_graph # TODO: hard-coding for MVP + self._forward_session = onnxruntime.InferenceSession(self._onnx_forward.SerializeToString()) + self._backward_session = onnxruntime.InferenceSession(self._onnx_backward.SerializeToString()) + # TODO: debug only + self._save_onnx_graph(self._onnx_forward, 'ortmodule_forward_mnist.onnx') + self._save_onnx_graph(self._onnx_backward, 'ortmodule_backward_mnist.onnx') - # TrainingParameters - # ort_parameters = onnxruntime.TrainingParameters() - # ort_parameters.loss_output_name = "loss" - # ort_parameters.use_mixed_precision = False - # ort_parameters.world_rank = 0 - # ort_parameters.world_size = 1 - # ort_parameters.gradient_accumulation_steps = 1 - # ort_parameters.allreduce_post_accumulation = False - # ort_parameters.deepspeed_zero_stage = 0 - # ort_parameters.enable_grad_norm_clip = False - # ort_parameters.set_gradients_as_graph_outputs = False - # ort_parameters.use_invertible_layernorm_grad = False - # ort_parameters.training_optimizer_name = "SGDOptimizer" - # ort_parameters.lr_params_feed_name = "Learning_Rate" - # ort_parameters.weights_to_train = trainable_params - # ort_parameters.optimizer_attributes_map = optimizer_attributes_map - # ort_parameters.optimizer_int_attributes_map = optimizer_int_attributes_map - - # # SessionOptions - # session_options = onnxruntime.SessionOptions() - # session_options.use_deterministic_compute = self.options.debug.deterministic_compute - # self.forward_session = onnxruntime.TrainingSession(self._onnx_forward.SerializeToString(), ort_parameters, session_options) - - self._save_onnx_graph(self._onnx_forward, 'forward_mnist.onnx') if not self._onnx_forward_initializers_desc: self._onnx_forward_initializers_desc = self._get_initializer_from_graph(self._onnx_forward) if not self._onnx_forward_inputs_desc: @@ -73,6 +55,7 @@ class ORTModule(torch.nn.Module): if not self._onnx_forward_outputs_desc: self._onnx_forward_outputs_desc = self._get_output_from_graph(self._onnx_forward) + # TODO: debug only print(f'Initializers: {self._onnx_forward_initializers_desc}') print(f'Inputs: {self._onnx_forward_inputs_desc}') print(f'Outpus: {self._onnx_forward_outputs_desc}') @@ -86,13 +69,12 @@ class ORTModule(torch.nn.Module): # Note: A potential optimization would be to detect which of inputs and weights # require a gradient. # intermediates, outputs = self._run_forward_graph(inputs) # inputs, weights) - # import pdb; pdb.set_trace() - outputs = self._run_forward_graph(*input, **kwargs) # inputs, weights) + outputs = self._run_forward_graph(*input, **kwargs) # inputs, weights) outputs = [torch.from_numpy(out).requires_grad_(True) for out in outputs] # TODO: Properly save intermediate tensors and remove them from model output - ctx.save_for_backward(outputs[1]) - outputs = [outputs[0]] + ctx.save_for_backward([(input, kwargs), outputs[1]]) + # outputs = [outputs[0]] # TODO: Properly support original module output format if len(outputs) == 1: @@ -100,42 +82,67 @@ class ORTModule(torch.nn.Module): return tuple(outputs) @staticmethod - def backward(ctx, grad_output): + def backward(ctx, *grad_output): print(f'_ORTModuleFunction.backward() was called') - ... - # intermediates = ctx.saved_tensors + input_and_kwargs, intermediates = ctx.saved_tensors # grad_inputs, grad_weights = self._run_backward_graph( # grad_output, intermediates) # return grad_inputs, grad_weights - return _ORTModuleFunction.apply(self._prepare_model_input(*input, **kwargs)) + return _ORTModuleFunction.apply(self._prepare_forward_input(*input, **kwargs)) - def _prepare_model_input(self, *input, **kwargs): + def _prepare_forward_input(self, *input, **kwargs): # Dictionary containing both inputs and initializers input_with_initializer = {} - # import pdb; pdb.set_trace() # Inputs - for idx, input_data in enumerate(self.forward_session.get_inputs()): - input_with_initializer.update({input_data.name : input[idx].cpu().numpy()}) + for idx, input_data in enumerate(self._forward_session.get_inputs()): + input_with_initializer.update({input_data.name: input[idx].cpu().numpy()}) # Initializers for idx, param in enumerate(self._original_module.named_parameters()): - input_with_initializer.update({param[0] : param[1].detach().numpy()}) + input_with_initializer.update({param[0]: param[1].detach().numpy()}) return input_with_initializer - def _run_forward_graph(self, data_with_initializer): #input, weights): - # import pdb; pdb.set_trace() - return self.forward_session.run(None, data_with_initializer) + def _prepare_backward_input(self, grad_output, intermediates, *inputs, **kwargs): + # Dictionary containing initializers + input_with_initializer = {} - def _run_backward_graph(self, grad_output, intermediates): + # User input + # TODO: How to determine which user input to feed to backward + for idx, input_data in enumerate(self._forward_session.get_inputs()): + input_with_initializer.update({input_data.name: inputs[idx].cpu().numpy()}) + + # Initializers + # TODO: How to determine which initializer (subset) to be used + for idx, param in enumerate(self._original_module.named_parameters()): + if param[0] == 'fc2.weight': + input_with_initializer.update({param[0]: param[1].detach().numpy()}) + + # Grad output + # TODO: How to determine grad_output name? + input_with_initializer.update({'probability_grad': grad_output.detach().numpy()}) + + # Intermediates + # TODO: How to determine intermediates name? + input_with_initializer.update({'7': intermediates.detach().numpy()}) + return input_with_initializer + + def _run_forward_graph(self, data_with_initializer): # input, weights): + print(f'_run_forward_graph was called...') + return self._forward_session.run(None, data_with_initializer) + + def _run_backward_graph(self, grad_output, intermediates, *inputs, **kwargs): # Use an InferenceSession to execute self.backward_graph. # Return gradient tensors for inputs and weights. - ... + print(f'_run_backward_graph was called...') + data = self._prepare_backward_input(grad_output, intermediates, *inputs, **kwargs) + return self._backward_session.run(None, data) @staticmethod def _get_forward_graph(module, module_input): + print(f'_get_forward_graph was called...') # TODO: Pytorch module must be exported to ONNX and splitted # Hard-coding with MNIST stub for MVP # Export torch.nn.Module to ONNX with initializers as input @@ -148,9 +155,10 @@ class ORTModule(torch.nn.Module): # export_params=True) # return onnx.load_model_from_string(f.getvalue()) return onnx.load('/home/thiagofc/mnist_onnx/mnist_with_training_forward_sliced.onnx') + # return onnx.load('/home/thiagofc/mnist_onnx/mnist_with_training.onnx') def _get_initializer_from_graph(self, graph): - # TODO: There is a tradefoo between memory footprint and total model export time + # TODO: There is a tradeoff between memory footprint and total model export time # Ideally we want to export the model using torch.onnx.export(.., export_params=False, keep_initializers_as_inputs=True) # to obtain an ONNX model with minimal size and initializers as input. # However, this results in (guessing) assuming only initializer's name end with '.weight' and '.bias'. @@ -168,7 +176,7 @@ class ORTModule(torch.nn.Module): # TODO: Dynamic shape is not being handled yet shape = initializer.dims dtype = _utils.dtype_onnx_to_torch(initializer.data_type) - initializers.append({'name' : name, 'shape' : shape, 'dtype' : dtype}) + initializers.append({'name': name, 'shape': shape, 'dtype': dtype}) return initializers def _get_input_from_graph(self, graph): @@ -182,7 +190,7 @@ class ORTModule(torch.nn.Module): # TODO: Dynamic shape is not being handled yet shape = [dim.dim_value for dim in elem.type.tensor_type.shape.dim] dtype = _utils.dtype_onnx_to_torch(elem.type.tensor_type.elem_type) - inputs.append({'name' : name, 'shape' : shape, 'dtype' : dtype}) + inputs.append({'name': name, 'shape': shape, 'dtype': dtype}) return inputs def _get_output_from_graph(self, graph): @@ -196,7 +204,7 @@ class ORTModule(torch.nn.Module): # TODO: Dynamic shape is not being handled yet shape = [dim.dim_value for dim in elem.type.tensor_type.shape.dim] dtype = _utils.dtype_onnx_to_torch(elem.type.tensor_type.elem_type) - outputs.append({'name' : name, 'shape' : shape, 'dtype' : dtype}) + outputs.append({'name': name, 'shape': shape, 'dtype': dtype}) return outputs @staticmethod @@ -226,18 +234,20 @@ class ORTModule(torch.nn.Module): @staticmethod def _build_gradient_graph(forward_graph): - # Invoke the C++ GradientBuilder implementation via pybind. + print(f'_build_gradient_graph was called...') + # TODO: Invoke the C++ GradientBuilder implementation via pybind. # Return an ONNX graph that contains the forward and backward nodes, which takes the # following inputs: # * Module inputs # * Module weights # * Gradients with respect to the module outputs # …and produces gradients with respect to the module inputs and weights. - ... + return onnx.load('/home/thiagofc/mnist_onnx/mnist_with_training_backward_sliced.onnx') @staticmethod def _split_forward_and_backward(gradient_graph): - # Split the result of _build_gradient_graph into two subgraphs: + print(f'_split_forward_and_backward was called...') + # TODO: Split the result of _build_gradient_graph into two subgraphs: # * A forward graph that takes module inputs and weights as input, and produces module # outputs and (“stashed”) intermediate tensors as output. # * A backward graph that takes intermediate tensors, module weights, and gradients diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic.py index 96f62f3828..9dd8d9b6e1 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic.py @@ -32,6 +32,7 @@ set_seed(seed) model = NeuralNet(input_size=784, hidden_size=500, num_classes=10) +print('Training MNIST on ORTModule....') model = ORTModule(model) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=lr) @@ -42,9 +43,31 @@ train_loader = torch.utils.data.DataLoader(datasets.MNIST('./data', train=True, transforms.Normalize((0.1307,), (0.3081,))])), batch_size=batch_size, shuffle=True) +# TODO: Get probability_grad from PyTorch Loss +probability_grad = torch.tensor([ +[0.36297542, 0.2297899, -0.10638658, 0.21579745, -0.12323117, -0.35163468, -0.16475351, -0.27790004, 0.20993066, 0.068910174], +[0.30177414, 0.4719398, -0.2290834, 0.61155605, -0.10533161, -0.068530589, -0.16963659, -0.034698304, 0.20859459, 0.071662053], +[0.26006302, 0.59704441, 0.2594507, 0.027483933, 0.17754407, -0.076404758, -0.15315992, -0.3511225, 0.096852496, -0.040248722], +[0.020109242, 0.47963268, 0.16444968, 0.28207836, 0.091335267, -0.34438723, -0.32664698, -0.04607122, 0.16735722, 0.28467956], +[-0.0067059044, 0.49364114, -0.023130134, 0.2933957, -0.12842584, -0.37883937, 0.083117418, -0.28517962, -0.021336049, -0.0058415309], +[-0.075187646, 0.24679491, 0.031593084, 0.59585023, -0.208859, -0.18786775, 0.18447922, -0.074010387, -0.056447648, -0.078843385], +[0.43958831, 0.53015679, -0.16698451, 0.3980948, 0.16000611, -0.016911259, -0.13209809, -0.10536471, 0.00073796883, 0.22187582], +[0.19641832, 0.47633961, 0.14354521, 0.49611267, -0.25266212, -0.28930596, -0.098222524, -0.17880601, 0.3030878, -0.086537011], +[0.16706356, 0.25445995, -0.36106035, 0.3932263, 0.020241318, -0.046459652, -0.30798167, 0.033364233, 0.10860923, 0.161856], +[0.076634176, 0.21363905, 0.14411786, 0.42425469, -0.36067143, -0.024277387, -0.23279551, -0.027842108, 0.11602029, 0.045313828], +[-0.067607164, 0.29514131, -0.21749593, 0.34894356, 0.10760085, -0.10467422, -0.39584625, 0.14010972, 0.21694142, 0.17883658], +[0.11919088, 0.17774329, -0.063672006, 0.31304225, 0.022851272, 0.00603014, -0.063586265, -0.11567068, 0.18024546, -0.044242512], +[0.28452805, 0.28950649, -0.030564137, 0.062676579, 0.037082255, -0.34579667, -0.18721311, -0.048553426, -0.047528304, -0.067283757], +[0.16541988, 0.6750235, 0.36633614, 0.12827933, -0.1848262, -0.12122689, 0.24612407, -0.22443134, 0.29384404, 0.029458519], +[0.022512322, -0.020067703, -0.035412017, 0.042415313, 0.01781881, -0.19647799, -0.019232273, -0.27665097, -0.085087284, -0.23508132], +[-0.056501552, 0.23281966, 0.012086541, 0.34509954, 0.096981436, -0.14569771, -0.24759589, 0.0071231984, 0.32205793, 0.027363759], +[-0.10276053, -0.15549006, 0.026301131, 0.067043148, -0.12606248, 0.042133313, -0.23401891, -0.16697425, -0.03425476, 0.14876992], +[0.20445672, 0.25619513, 0.16442557, 0.077375375, 0.13566223, -0.099527359, -0.12576742, -0.45158958, 0.32187107, 0.092045955], +[0.34017974, -0.066395164, 0.20674077, 0.16103405, -0.27109221, -0.24286765, -0.14018115, -0.0068955906, 0.17458764, -0.072009444], +[-0.081807368, 0.30574301, -0.15613964, 0.33026001, -0.12889105, -0.053762466, 0.036609523, -0.16667747, 0.12113887, -0.10802352], +]) # Training Loop -print('Training MNIST on ORTModule....') loss = float('inf') for iteration, (data, target) in enumerate(train_loader): if iteration == 1: @@ -53,11 +76,27 @@ for iteration, (data, target) in enumerate(train_loader): data = data.reshape(data.shape[0], -1) optimizer.zero_grad() - probability = model(data) - print(f'Output from forward has shape {probability.size()}: {probability}') - # import pdb; pdb.set_trace() + probability, intermediates = model(data) + print(f'Output from forward has shape {probability.size()}') loss = criterion(probability, target) - # loss.backward() + + # Fake backward call to test backprop graph + fc1_bias_grad, fc1_weight_grad, fc2_weight_grad, fc2_bias_grad = model._run_backward_graph(probability_grad, intermediates, data) + fc1_bias_grad = torch.nn.Parameter(torch.from_numpy(fc1_bias_grad)) + fc2_bias_grad = torch.nn.Parameter(torch.from_numpy(fc2_bias_grad)) + fc1_weight_grad = torch.nn.Parameter(torch.from_numpy(fc1_weight_grad)) + fc2_weight_grad = torch.nn.Parameter(torch.from_numpy(fc2_weight_grad)) + model._original_module.fc1.bias = fc1_bias_grad + model._original_module.fc1.weight = fc1_weight_grad + model._original_module.fc2.bias = fc2_bias_grad + model._original_module.fc2.weight = fc2_weight_grad + + print(f'Output from backaward has the following shapes after update:') + print(f'fc1_bias_grad={fc1_bias_grad.size()}') + print(f'fc2_bias_grad={fc2_bias_grad.size()}') + print(f'fc1_weight_grad={fc1_weight_grad.size()}') + print(f'fc2_weight_grad={fc2_weight_grad.size()}') + # loss.backward(target) # optimizer.step() if iteration == 0: diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic_transform_model.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic_transform_model.py index 7fbc6a2d3e..0965b5ad06 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic_transform_model.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_basic_transform_model.py @@ -1,4 +1,5 @@ # coding=utf8 +import copy import sys import onnx from onnx import helper, shape_inference @@ -14,11 +15,11 @@ input_model_name = sys.argv[1] output_forward_model_name = input_model_name[:-5] + '_forward_sliced.onnx' output_backward_model_name = input_model_name[:-5] + '_backward_sliced.onnx' -def add_model_input_from_initializer(model, initializer, docstring=None): +def add_input_from_initializer(model, initializer, docstring=None): new_input = onnx.helper.make_tensor_value_info(initializer.name, initializer.data_type, initializer.dims, docstring) model.graph.input.append(new_input) -def add_model_input(model, name, data_type, dims, docstring=None): +def add_input(model, name, data_type = None, dims = None, docstring = None): new_input = onnx.helper.make_tensor_value_info(name, data_type, dims, docstring) model.graph.input.append(new_input) @@ -61,73 +62,45 @@ def find_node(model, name): model = onnx.load(input_model_name) # Remove model inputs -# They are: label -node = find_model_input(model, 'label') -model.graph.input.remove(node) +nodes = ['label'] +for node in nodes: + node = find_model_input(model, node) + model.graph.input.remove(node) # Remove model outputs -# They are: loss -node = find_model_output(model, 'loss') -model.graph.output.remove(node) -node = find_model_output(model, 'fc1.bias_grad') -model.graph.output.remove(node) -node = find_model_output(model, 'fc1.weight_grad') -model.graph.output.remove(node) -node = find_model_output(model, 'fc2.bias_grad') -model.graph.output.remove(node) -node = find_model_output(model, 'fc2.weight_grad') -model.graph.output.remove(node) +nodes = ['loss', 'fc1.bias_grad', 'fc1.weight_grad', 'fc2.bias_grad', 'fc2.weight_grad'] +for node in nodes: + node = find_model_output(model, node) + model.graph.output.remove(node) # Add input with same name, type and shape as the initializers # They are: [fc1.bias, fc1.weight, fc2.bias, fc2.weight] -node = find_initializer(model, 'fc1.bias') -add_model_input_from_initializer(model, node, 'thiagofc: add fc1.bias as model input') -node = find_initializer(model, 'fc1.weight') -add_model_input_from_initializer(model, node, 'thiagofc: add fc1.weight as model input') -node = find_initializer(model, 'fc2.bias') -add_model_input_from_initializer(model, node, 'thiagofc: add fc2.bias as model input') -node = find_initializer(model, 'fc2.weight') -add_model_input_from_initializer(model, node, 'thiagofc: add fc2.weight as model input') +forward_initializer_names = ['fc1.bias', 'fc1.weight', 'fc2.bias', 'fc2.weight'] +forward_initializer = {} +for node in forward_initializer_names: + node = find_initializer(model, node) + add_input_from_initializer(model, node, f'thiagofc: add {node.name} as model input') + forward_initializer.update({node.name : copy.deepcopy(node)}) # Remove initializers from model -# They are: [fc1.bias, fc1.weight, fc2.bias, fc2.weight] # TODO: Do this when we are able to distinguish inputs from initializers -# model.graph.initializer.remove(node) -# model.graph.initializer.remove(node) -# model.graph.initializer.remove(node) -# model.graph.initializer.remove(node) +# for node in forward_initializer_names: +# node = find_initializer(model, init) +# model.graph.initializer.remove(node) # Remove backward-related initializers -# They are: [loss_grad, ZeroConstant] -node = find_initializer(model, 'loss_grad') -model.graph.initializer.remove(node) -node = find_initializer(model, 'ZeroConstant') -model.graph.initializer.remove(node) +nodes = ['loss_grad', 'ZeroConstant'] +for node in nodes: + node = find_initializer(model, node) + model.graph.initializer.remove(node) # Remove OPs -# They are: [SoftmaxCrossEntropyLoss_3, SoftmaxCrossEntropyLoss_3_Grad/SoftmaxCrossEntropyLossGrad_0, -# Gemm_2_Grad/ReduceSum_3, Gemm_2_Grad/Identity_4, Gemm_2_Grad/Gemm_2, Gemm_2_Grad/Gemm_1, -# Relu_1_Grad/ReluGrad_0, Gemm_0_Grad/Gemm_1, Gemm_0_Grad/ReduceSum_2, Gemm_0_Grad/Identity_3] -node = find_node(model, 'SoftmaxCrossEntropyLoss_3') -model.graph.node.remove(node) -node = find_node(model, 'SoftmaxCrossEntropyLoss_3_Grad/SoftmaxCrossEntropyLossGrad_0') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_2_Grad/ReduceSum_3') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_2_Grad/Identity_4') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_2_Grad/Gemm_2') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_2_Grad/Gemm_1') -model.graph.node.remove(node) -node = find_node(model, 'Relu_1_Grad/ReluGrad_0') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_0_Grad/Gemm_1') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_0_Grad/ReduceSum_2') -model.graph.node.remove(node) -node = find_node(model, 'Gemm_0_Grad/Identity_3') -model.graph.node.remove(node) +nodes = ['SoftmaxCrossEntropyLoss_3', 'SoftmaxCrossEntropyLoss_3_Grad/SoftmaxCrossEntropyLossGrad_0', + 'Gemm_2_Grad/ReduceSum_3', 'Gemm_2_Grad/Identity_4', 'Gemm_2_Grad/Gemm_2', 'Gemm_2_Grad/Gemm_1', + 'Relu_1_Grad/ReluGrad_0', 'Gemm_0_Grad/Gemm_1', 'Gemm_0_Grad/ReduceSum_2', 'Gemm_0_Grad/Identity_3'] +for node in nodes: + node = find_node(model, node) + model.graph.node.remove(node) # Add new outputs: # They are: 7 @@ -138,11 +111,38 @@ with open(output_forward_model_name, "wb") as f: ############################################################################### -# FORWARD PASS GRAPH ########################################################## +# BACKWARD PASS GRAPH ########################################################## ############################################################################### model = onnx.load(input_model_name) +# Add new inputs: +# TODO: Should we specify types here? ORT graph doesn't have that info available, but ONNX API needs it +add_input_from_initializer(model, forward_initializer['fc2.weight'], 'thiagofc: add fc2.weight as model input') +add_input(model, '7', 1, None, 'thiagofc: add 7 as model input') +add_input(model, 'probability_grad', 1, None, 'thiagofc: add probability_grad as model input') +# Remove model inputs +# They are: label +node = find_model_input(model, 'label') +model.graph.input.remove(node) + +# Remove model outputs +nodes = ['loss', 'probability'] +for node in nodes: + node = find_model_output(model, node) + model.graph.output.remove(node) + +# Remove OP nodes from forward pass +nodes = ['Gemm_0', 'Relu_1', 'Gemm_2', 'SoftmaxCrossEntropyLoss_3', 'SoftmaxCrossEntropyLoss_3_Grad/SoftmaxCrossEntropyLossGrad_0'] +for node in nodes: + node = find_node(model, node) + model.graph.node.remove(node) + +# Remove initializers +forward_initializer_names.extend(['loss_grad']) +for node in forward_initializer_names: + node = find_initializer(model, node) + model.graph.initializer.remove(node) with open(output_backward_model_name, "wb") as f: f.write(model.SerializeToString()) diff --git a/samples/python/mnist/mnist_training.py b/samples/python/mnist/ort_mnist.py similarity index 78% rename from samples/python/mnist/mnist_training.py rename to samples/python/mnist/ort_mnist.py index 3d23987e60..428a252a14 100644 --- a/samples/python/mnist/mnist_training.py +++ b/samples/python/mnist/ort_mnist.py @@ -1,6 +1,14 @@ # This code is from https://github.com/pytorch/examples/blob/master/mnist/main.py # with modification to do training using onnxruntime as backend on cuda device. +# To print nodes from ORT backend +# Add --cmake_extra_defines onnxruntime_DEBUG_NODE_INPUTS_OUTPUTS=1 to build.sh +# export ORT_DEBUG_NODE_IO_NAME_FILTER="SoftmaxCrossEntropyLoss_3_Grad/SoftmaxCrossEntropyLossGrad_0" +# export ORT_DEBUG_NODE_IO_NAME_FILTER="SoftmaxCrossEntropyLoss_3" +# export ORT_DEBUG_NODE_IO_DUMP_INPUT_DATA=1 +# export ORT_DEBUG_NODE_IO_DUMP_OUTPUT_DATA=1 +# See https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/framework/debug_node_inputs_outputs_utils.h + import argparse import torch import torch.nn as nn @@ -33,20 +41,25 @@ def mnist_model_description(): 'outputs': [('loss', [], True), ('probability', ['batch', 10])]} - def my_loss(x, target): return F.nll_loss(F.log_softmax(x, dim=1), target) - # Helpers -def train_with_trainer(log_interval, trainer, device, train_loader, epoch): +def train(log_interval, trainer, device, train_loader, epoch): for batch_idx, (data, target) in enumerate(train_loader): # Fetch data data, target = data.to(device), target.to(device) data = data.reshape(data.shape[0], -1) # Train step - loss, _ = trainer.train_step(data, target) + loss, prob = trainer.train_step(data, target) + + if batch_idx == 0: + # trainer.save_as_onnx('/home/thiagofc/mnist_onnx/pytorch_as_onnx.onnx') + # import pdb; pdb.set_trace() + pass + else: + break # Stats if batch_idx % log_interval == 0: @@ -55,16 +68,14 @@ def train_with_trainer(log_interval, trainer, device, train_loader, epoch): 100. * batch_idx / len(train_loader), loss)) -def test_with_trainer(trainer, device, test_loader): +def test(trainer, device, test_loader): test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: - # Fetch data data, target = data.to(device), target.to(device) data = data.reshape(data.shape[0], -1) - # Eval step # Using fetches around without eval_step to not pass 'target' as input trainer._train_step_info.fetches = ['probability'] output = F.log_softmax(trainer.eval_step(data), dim=1) @@ -74,7 +85,9 @@ def test_with_trainer(trainer, device, test_loader): test_loss += F.nll_loss(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() + test_loss /= len(test_loader.dataset) + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) @@ -82,13 +95,13 @@ def test_with_trainer(trainer, device, test_loader): def main(): # Training settings - parser = argparse.ArgumentParser(description='MNIST Example') - parser.add_argument('--batch-size', type=int, default=64, metavar='N', - help='input batch size for training (default: 64)') + parser = argparse.ArgumentParser(description='ONNX Runtime MNIST Example') + parser.add_argument('--batch-size', type=int, default=20, metavar='N', + help='input batch size for training (default: 20)') parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)') - parser.add_argument('--epochs', type=int, default=10, metavar='N', - help='number of epochs to train (default: 10)') + parser.add_argument('--epochs', type=int, default=1, metavar='N', + help='number of epochs to train (default: 1)') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--no-cuda', action='store_true', default=False, @@ -125,6 +138,10 @@ def main(): model_desc = mnist_model_description() optim_config = optim.SGDConfig(lr=args.lr) opts = ORTTrainerOptions({'device': {'id': device}}) + + # import onnx + # model = onnx.load('/home/thiagofc/mnist_onnx/mnist_with_training_probability_grad.onnx') + # my_loss=None trainer = ORTTrainer(model, model_desc, optim_config, @@ -133,9 +150,8 @@ def main(): # Train loop for epoch in range(1, args.epochs + 1): - train_with_trainer(args.log_interval, trainer, - device, train_loader, epoch) - test_with_trainer(trainer, device, test_loader) + train(args.log_interval, trainer, device, train_loader, epoch) + # test(trainer, device, test_loader) if __name__ == '__main__': diff --git a/samples/python/mnist/pytorch_mnist.py b/samples/python/mnist/pytorch_mnist.py new file mode 100644 index 0000000000..4f2b5f0927 --- /dev/null +++ b/samples/python/mnist/pytorch_mnist.py @@ -0,0 +1,119 @@ +import argparse +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torchvision import datasets, transforms + + +# Pytorch model +class NeuralNet(nn.Module): + def __init__(self, input_size, hidden_size, num_classes): + super(NeuralNet, self).__init__() + self.fc1 = nn.Linear(input_size, hidden_size) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(hidden_size, num_classes) + + def forward(self, input1): + out = self.fc1(input1) + out = self.relu(out) + out = self.fc2(out) + return out + + +def my_loss(x, target, is_train=True): + if is_train: + return F.nll_loss(F.log_softmax(x, dim=1), target) + else: + return F.nll_loss(F.log_softmax(x, dim=1), target, reduction='sum') + +# Helpers +def train(args, model, device, train_loader, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(train_loader): + data, target = data.to(device), target.to(device) + data = data.reshape(data.shape[0], -1) + optimizer.zero_grad() + output = model(data) + loss = my_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % args.log_interval == 0: + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( + epoch, batch_idx * len(data), len(train_loader.dataset), + 100. * batch_idx / len(train_loader), loss.item())) + + +def test(model, device, test_loader): + model.eval() + test_loss = 0 + correct = 0 + with torch.no_grad(): + for data, target in test_loader: + data, target = data.to(device), target.to(device) + data = data.reshape(data.shape[0], -1) + output = model(data) + # Stats + test_loss += my_loss(output, target, False).item() + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + + test_loss /= len(test_loader.dataset) + + print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( + test_loss, correct, len(test_loader.dataset), + 100. * correct / len(test_loader.dataset))) + + +def main(): + # Training settings + parser = argparse.ArgumentParser(description='PyTorch MNIST Example') + parser.add_argument('--batch-size', type=int, default=64, metavar='N', + help='input batch size for training (default: 64)') + parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', + help='input batch size for testing (default: 1000)') + parser.add_argument('--epochs', type=int, default=1, metavar='N', + help='number of epochs to train (default: 1)') + parser.add_argument('--lr', type=float, default=0.01, metavar='LR', + help='learning rate (default: 0.01)') + parser.add_argument('--no-cuda', action='store_true', default=False, + help='disables CUDA training') + parser.add_argument('--seed', type=int, default=1, metavar='S', + help='random seed (default: 1)') + parser.add_argument('--log-interval', type=int, default=10, metavar='N', + help='how many batches to wait before logging training status') + + # Basic setup + args = parser.parse_args() + if not args.no_cuda and torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + torch.manual_seed(args.seed) + + # Data loader + train_loader = torch.utils.data.DataLoader( + datasets.MNIST('./data', train=True, download=True, + transform=transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.1307,), (0.3081,)) + ])), + batch_size=args.batch_size, shuffle=True) + test_loader = torch.utils.data.DataLoader( + datasets.MNIST('./data', train=False, transform=transforms.Compose([ + transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])), + batch_size=args.test_batch_size, shuffle=True) + + # Modeling + model = NeuralNet(784, 500, 10).to(device) + optimizer = optim.SGD(model.parameters(), lr=args.lr) + + # Train loop + for epoch in range(1, args.epochs + 1): + train(args, model, device, train_loader, optimizer, epoch) + test(model, device, test_loader) + optimizer.step() + + +if __name__ == '__main__': + main()