mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-05 04:17:53 +00:00
Basic plumbing for backward pass. Not fully working
This commit is contained in:
Thiago Crepaldi
parent
77cefcd6c2
commit
e71e08851a
5 changed files with 316 additions and 132 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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__':
|
||||
119
samples/python/mnist/pytorch_mnist.py
Normal file
119
samples/python/mnist/pytorch_mnist.py
Normal file
|
|
@ -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()
|
||||
Loading…
Reference in a new issue