mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Add ONNX Training Post-Passes to Front-End - Cont (#4041)
* Add ONNX postpasses * add flag + add bert test from onnx file * address PR comments * fix typo * fix rebase * address comments * Fix test failures * add new pass for expand for new pt version, add comments * fix rebase Co-authored-by: lahaidar <lahaidar@OrtDevTest2v100.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Bowen Bao
GitHub
parent
0b5bbb16b8
commit
b08771f00e
6 changed files with 704 additions and 24 deletions
|
|
@ -727,7 +727,7 @@ class TestOrtTrainer(unittest.TestCase):
|
|||
model_desc = ModelDescription([input_desc, label_desc], [loss_desc, output_desc])
|
||||
def loss_fn(x, label):
|
||||
return F.nll_loss(F.log_softmax(x, dim=1), label)
|
||||
|
||||
|
||||
def get_lr_this_step(global_step):
|
||||
learningRate = 0.02
|
||||
return torch.tensor([learningRate])
|
||||
|
|
|
|||
|
|
@ -8,14 +8,15 @@ from onnx import helper
|
|||
import torch
|
||||
import torch.nn
|
||||
import torch.onnx
|
||||
import onnxruntime as ort
|
||||
import onnxruntime.capi.postprocess as postprocess
|
||||
from distutils.version import LooseVersion
|
||||
import warnings
|
||||
|
||||
import onnxruntime as ort
|
||||
from .checkpointing_utils import list_checkpoint_files, get_checkpoint_name, CombineZeroCheckpoint
|
||||
import onnxruntime.capi.pt_patch
|
||||
|
||||
DEFAULT_OPSET_VERSION = 10
|
||||
DEFAULT_OPSET_VERSION = 12
|
||||
|
||||
class IODescription():
|
||||
def __init__(self, name, shape, dtype=None, num_classes=None):
|
||||
|
|
@ -204,7 +205,7 @@ def wrap_for_input_match(model, loss_fn, input_names):
|
|||
if len(sig_loss.parameters) != 2:
|
||||
raise RuntimeError("loss function should take two arguments - predict and label.")
|
||||
|
||||
# label shall be the second input to loss_fn.
|
||||
# label shall be the second input to loss_fn.
|
||||
ordered_list_keys = [*ordered_list_keys, list(sig_loss.parameters.keys())[1]]
|
||||
|
||||
class model_loss_cls(torch.nn.Module):
|
||||
|
|
@ -274,7 +275,7 @@ def wrap_for_input_match(model, loss_fn, input_names):
|
|||
|
||||
return model
|
||||
|
||||
def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, opset_version=DEFAULT_OPSET_VERSION):
|
||||
def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, opset_version=DEFAULT_OPSET_VERSION, _enable_internal_postprocess=True):
|
||||
# example: {input0:{0:'batch'}, input1:{0:'batch'}}
|
||||
dynamic_axes = {}
|
||||
for input in model_desc.inputs_:
|
||||
|
|
@ -323,17 +324,20 @@ def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, op
|
|||
|
||||
# Other export options to use(this is for backward compatibility).
|
||||
other_export_options = {}
|
||||
other_export_options['training'] = True
|
||||
|
||||
# This option was added after 1.4 release.
|
||||
if LooseVersion(torch.__version__) > LooseVersion('1.4.0'):
|
||||
other_export_options['enable_onnx_checker'] = False
|
||||
# This option was added after 1.6 release.
|
||||
if LooseVersion(torch.__version__) >= LooseVersion('1.6.0'):
|
||||
other_export_options['training'] = torch.onnx.TrainingMode.TRAINING
|
||||
|
||||
torch.onnx._export(model, tuple(sample_inputs), f,
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
opset_version=opset_version,
|
||||
dynamic_axes=dynamic_axes,
|
||||
training=True,
|
||||
_retain_param_name=True,
|
||||
example_outputs=tuple(sample_outputs),
|
||||
do_constant_folding=False,
|
||||
|
|
@ -360,7 +364,8 @@ def convert_model_loss_fn_to_onnx(model, loss_fn, model_desc, device, inputs, op
|
|||
"Initializer names do not match between PyTorch model and ONNX model, " \
|
||||
"please report a bug to ONNX Runtime."
|
||||
|
||||
onnx_model = FuseSofmaxNLLToSoftmaxCE(onnx_model)
|
||||
if _enable_internal_postprocess:
|
||||
onnx_model = postprocess.run_postprocess(onnx_model)
|
||||
|
||||
return onnx_model
|
||||
|
||||
|
|
@ -518,10 +523,11 @@ def load_checkpoint(model, checkpoint_dir, checkpoint_prefix="ORT_checkpoint", s
|
|||
class ORTTrainer():
|
||||
|
||||
def __init__(self, model, loss_fn, model_desc, training_optimizer_name, map_optimizer_attributes,
|
||||
learning_rate_description, device, gradient_accumulation_steps=1, postprocess_model=None,
|
||||
learning_rate_description, device, gradient_accumulation_steps=1,
|
||||
world_rank=0, world_size=1, use_mixed_precision=False, allreduce_post_accumulation=False,
|
||||
global_step=0, get_lr_this_step=None, loss_scaler=None, partition_optimizer=False,
|
||||
enable_grad_norm_clip=True, frozen_weights=[], _opset_version=DEFAULT_OPSET_VERSION):
|
||||
enable_grad_norm_clip=True, frozen_weights=[], _opset_version=DEFAULT_OPSET_VERSION,
|
||||
_enable_internal_postprocess=True, _extra_postprocess=None):
|
||||
super(ORTTrainer, self).__init__()
|
||||
"""
|
||||
Initialize ORTTrainer.
|
||||
|
|
@ -586,6 +592,10 @@ class ORTTrainer():
|
|||
Defaults to True.
|
||||
frozen_weights: list of model parameters to be frozen (not trained).
|
||||
Defaults to [].
|
||||
_enable_internal_postprocess: whether to run or not the internal postprocesses.
|
||||
Defaults to True
|
||||
_extra_postprocess: a callable to postprocess the ONNX model that is converted from PyTorch.
|
||||
Defaults to None
|
||||
"""
|
||||
warnings.warn('DISCLAIMER: This is an early version of an experimental training API and it is subject to change. DO NOT create production applications with it')
|
||||
self.is_train = True
|
||||
|
|
@ -621,7 +631,7 @@ class ORTTrainer():
|
|||
# we use self.global_step_ to count optimizations being performed.
|
||||
# it is used to calculate learning rate if self.get_lr_this_step_ is provided.
|
||||
self.global_step_ = global_step
|
||||
self.post_process_model_fn_ = postprocess_model
|
||||
self._extra_postprocess = _extra_postprocess
|
||||
self.get_lr_this_step_ = get_lr_this_step
|
||||
self.loss_scaler_ = loss_scaler
|
||||
|
||||
|
|
@ -636,6 +646,7 @@ class ORTTrainer():
|
|||
self.frozen_weights_ = frozen_weights
|
||||
self.opset_version_ = _opset_version
|
||||
self.state_dict_ = None
|
||||
self._enable_internal_postprocess = _enable_internal_postprocess
|
||||
|
||||
# use this special string to workaround a corner case that external loss_scale is passed into train_step as kwargs.
|
||||
# see prepare_input_and_fetches for more details.
|
||||
|
|
@ -647,6 +658,12 @@ class ORTTrainer():
|
|||
if self.onnx_model_ is None:
|
||||
return
|
||||
|
||||
if self._enable_internal_postprocess:
|
||||
self._onnx_model_ = postprocess.run_postprocess(self.onnx_model_)
|
||||
|
||||
if self._extra_postprocess:
|
||||
self._extra_postprocess(self.onnx_model_)
|
||||
|
||||
self._verify_fully_optimized_model(self.onnx_model_)
|
||||
self.session, self.train_io_binding, self.eval_io_binding, self.output_name, _, self.output_types = \
|
||||
create_ort_training_session_with_optimizer(
|
||||
|
|
@ -702,10 +719,10 @@ class ORTTrainer():
|
|||
torch_buffers = list(dict(self.torch_model_.named_buffers()).keys())
|
||||
self.frozen_weights_ = self.frozen_weights_ + torch_buffers
|
||||
self.onnx_model_ = convert_model_loss_fn_to_onnx(
|
||||
self.torch_model_, self.loss_fn_, self.model_desc_, torch.device('cpu'), inputs, opset_version=self.opset_version_)
|
||||
self.torch_model_, self.loss_fn_, self.model_desc_, torch.device('cpu'), inputs, opset_version=self.opset_version_, _enable_internal_postprocess=self._enable_internal_postprocess)
|
||||
|
||||
if self.post_process_model_fn_:
|
||||
self.post_process_model_fn_(self.onnx_model_)
|
||||
if self._extra_postprocess:
|
||||
self._extra_postprocess(self.onnx_model_)
|
||||
|
||||
self._init_session()
|
||||
|
||||
|
|
|
|||
455
orttraining/orttraining/python/postprocess.py
Normal file
455
orttraining/orttraining/python/postprocess.py
Normal file
|
|
@ -0,0 +1,455 @@
|
|||
import sys
|
||||
import os.path
|
||||
from onnx import *
|
||||
import onnx
|
||||
import numpy as np
|
||||
import struct
|
||||
|
||||
from onnx import helper
|
||||
from onnx import numpy_helper
|
||||
|
||||
|
||||
def run_postprocess(model):
|
||||
# this post pass is not required for pytorch >= 1.5
|
||||
# where add_node_name in torch.onnx.export is default to True
|
||||
model = add_name(model)
|
||||
|
||||
# this post pass is not required for pytorch > 1.6
|
||||
model = fuse_softmaxNLL_to_softmaxCE(model)
|
||||
|
||||
model = layer_norm_transform(model)
|
||||
model = fix_expand_shape(model)
|
||||
model = fix_expand_shape_pt_1_5(model)
|
||||
return model
|
||||
|
||||
def find_input_node(model, arg):
|
||||
result = []
|
||||
for node in model.graph.node:
|
||||
for output in node.output:
|
||||
if output == arg:
|
||||
result.append(node)
|
||||
return result[0] if len(result)== 1 else None
|
||||
|
||||
def find_output_node(model, arg):
|
||||
result = []
|
||||
for node in model.graph.node:
|
||||
for input in node.input:
|
||||
if input == arg:
|
||||
result.append(node)
|
||||
return result[0] if len(result) == 1 else result
|
||||
|
||||
def add_name(model):
|
||||
i = 0
|
||||
for node in model.graph.node:
|
||||
node.name = '%s_%d' %(node.op_type, i)
|
||||
i += 1
|
||||
return model
|
||||
|
||||
# Expand Shape PostProcess
|
||||
|
||||
def fix_expand_shape(model):
|
||||
expand_nodes = [n for n in model.graph.node if n.op_type == 'Expand']
|
||||
model_inputs_names = [i.name for i in model.graph.input]
|
||||
|
||||
for expand_node in expand_nodes:
|
||||
shape = find_input_node(model, expand_node.input[1])
|
||||
if shape.op_type == 'Shape':
|
||||
# an expand subgraph
|
||||
# Input Input2
|
||||
# | |
|
||||
# | Shape
|
||||
# | |
|
||||
# |__ __|
|
||||
# | |
|
||||
# Expand
|
||||
# |
|
||||
# output
|
||||
#
|
||||
# Only if Input2 is one of the model inputs, assign Input2's shape to output of expand.
|
||||
shape_input_name = shape.input[0]
|
||||
if shape_input_name in model_inputs_names:
|
||||
index = model_inputs_names.index(shape_input_name)
|
||||
expand_out = model.graph.value_info.add()
|
||||
expand_out.name = expand_node.output[0]
|
||||
expand_out.type.CopyFrom(model.graph.input[index].type)
|
||||
return model
|
||||
|
||||
def fix_expand_shape_pt_1_5(model):
|
||||
# expand subgraph
|
||||
# Constant
|
||||
# +
|
||||
# ConstantOfShape
|
||||
# | + |
|
||||
# | + |
|
||||
# (Reshape subgraph) Mul |
|
||||
# |___ _________| |
|
||||
# + | | |
|
||||
# + Equal |
|
||||
# +++++|++++++++++++++|++
|
||||
# |____________ | +
|
||||
# | | +
|
||||
# (subgraph) Where
|
||||
# | |
|
||||
# |_____ ___________|
|
||||
# | |
|
||||
# Expand
|
||||
# |
|
||||
# output
|
||||
#
|
||||
# where the Reshape subgraph is
|
||||
#
|
||||
# Input
|
||||
# | |
|
||||
# | |___________________
|
||||
# | |
|
||||
# Shape Constant Shape Constant
|
||||
# | ______| | ______|
|
||||
# | | | |
|
||||
# Gather Gather
|
||||
# | |
|
||||
# Unsqueeze Unsqueeze
|
||||
# | |
|
||||
# | ..Number of dims.. |
|
||||
# | _________________|
|
||||
# |...|
|
||||
# Concat Constant
|
||||
# | |
|
||||
# |______ __________________|
|
||||
# | |
|
||||
# Reshape
|
||||
# |
|
||||
# output
|
||||
#
|
||||
# This pass will copy Input's shape to the output of Expand.
|
||||
expand_nodes = [n for n in model.graph.node if n.op_type == 'Expand']
|
||||
model_inputs_names = [i.name for i in model.graph.input]
|
||||
|
||||
for expand_node in expand_nodes:
|
||||
n_where = find_input_node(model, expand_node.input[1])
|
||||
if n_where.op_type != 'Where':
|
||||
continue
|
||||
|
||||
n_equal = find_input_node(model, n_where.input[0])
|
||||
n_cos = find_input_node(model, n_where.input[1])
|
||||
n_reshape = find_input_node(model, n_where.input[2])
|
||||
|
||||
if n_equal.op_type != 'Equal' or n_cos.op_type != 'ConstantOfShape' or n_reshape.op_type != 'Reshape':
|
||||
continue
|
||||
|
||||
n_reshape_e = find_input_node(model, n_equal.input[0])
|
||||
n_mul = find_input_node(model, n_equal.input[1])
|
||||
if n_reshape_e != n_reshape or n_mul.op_type != 'Mul':
|
||||
continue
|
||||
|
||||
n_cos_m = find_input_node(model, n_mul.input[0])
|
||||
n_constant = find_input_node(model, n_mul.input[1])
|
||||
if n_cos_m != n_cos or n_constant.op_type != 'Constant':
|
||||
continue
|
||||
|
||||
n_concat = find_input_node(model, n_reshape.input[0])
|
||||
n_constant_r = find_input_node(model, n_reshape.input[1])
|
||||
if n_concat.op_type != 'Concat' or n_constant_r.op_type != 'Constant':
|
||||
continue
|
||||
|
||||
n_input_candidates = []
|
||||
for concat_in in n_concat.input:
|
||||
n_unsqueeze = find_input_node(model, concat_in)
|
||||
if n_unsqueeze.op_type != 'Unsqueeze':
|
||||
break
|
||||
n_gather = find_input_node(model, n_unsqueeze.input[0])
|
||||
if n_gather.op_type != 'Gather':
|
||||
break
|
||||
n_shape = find_input_node(model, n_gather.input[0])
|
||||
n_constant_g = find_input_node(model, n_gather.input[1])
|
||||
if n_shape.op_type != 'Shape' or n_constant_g.op_type != 'Constant':
|
||||
break
|
||||
n_input = n_shape.input[0]
|
||||
if not n_input in model_inputs_names:
|
||||
break
|
||||
n_input_candidates.append(n_input)
|
||||
|
||||
if not n_input_candidates or not all(elem == n_input_candidates[0] for elem in n_input_candidates):
|
||||
continue
|
||||
|
||||
index = model_inputs_names.index(n_input_candidates[0])
|
||||
expand_out = model.graph.value_info.add()
|
||||
expand_out.name = expand_node.output[0]
|
||||
expand_out.type.CopyFrom(model.graph.input[index].type)
|
||||
return model
|
||||
|
||||
# LayerNorm PostProcess
|
||||
|
||||
def find_nodes(graph, op_type):
|
||||
nodes = []
|
||||
for node in graph.node:
|
||||
if node.op_type == op_type:
|
||||
nodes.append(node)
|
||||
return nodes
|
||||
|
||||
def is_type(node, op_type):
|
||||
if node is None or isinstance(node, list):
|
||||
return False
|
||||
return node.op_type == op_type
|
||||
|
||||
def add_const(model, name, output, t_value = None, f_value = None):
|
||||
const_node = model.graph.node.add()
|
||||
const_node.op_type = 'Constant'
|
||||
const_node.name = name
|
||||
const_node.output.extend([output])
|
||||
attr = const_node.attribute.add()
|
||||
attr.name = 'value'
|
||||
if t_value is not None:
|
||||
attr.type = 4
|
||||
attr.t.CopyFrom(t_value)
|
||||
else:
|
||||
attr.type = 1
|
||||
attr.f = f_value
|
||||
return const_node
|
||||
|
||||
def layer_norm_transform(model):
|
||||
# Converting below subgraph
|
||||
#
|
||||
# input
|
||||
# |
|
||||
# ReduceMean
|
||||
# __|_____
|
||||
# | |
|
||||
# Sub Sub
|
||||
# | |
|
||||
# | (optional) Cast
|
||||
# | |
|
||||
# | Pow
|
||||
# | |
|
||||
# | (optional) Cast
|
||||
# | |
|
||||
# | ReduceMean
|
||||
# | |
|
||||
# | Add
|
||||
# | |
|
||||
# |__ __Sqrt
|
||||
# | |
|
||||
# Div (weight)
|
||||
# | |
|
||||
# | _____|
|
||||
# | |
|
||||
# Mul (bias)
|
||||
# | |
|
||||
# | _____|
|
||||
# | |
|
||||
# Add
|
||||
# |
|
||||
# output
|
||||
#
|
||||
# to the below subgraph
|
||||
#
|
||||
# input (weight) (bias)
|
||||
# | | |
|
||||
# | _______| |
|
||||
# | | ________________|
|
||||
# | | |
|
||||
# LayerNormalization
|
||||
# |
|
||||
# output
|
||||
graph = model.graph
|
||||
|
||||
nodes_ReduceMean = find_nodes(graph, "ReduceMean")
|
||||
|
||||
id = 0
|
||||
layer_norm_nodes = []
|
||||
remove_nodes = []
|
||||
for reduce_mean in nodes_ReduceMean:
|
||||
# check that reduce_mean output is Sub
|
||||
sub = find_output_node(model, reduce_mean.output[0])
|
||||
if not is_type(sub, "Sub"):
|
||||
continue
|
||||
|
||||
# check that sub output[0] is Div and output[1] is Pow
|
||||
pow, div = find_output_node(model, sub.output[0])
|
||||
if is_type(pow, "Cast"):
|
||||
# During an update in PyTorch, Cast nodes are inserted between Sub and Pow.
|
||||
remove_nodes += [pow]
|
||||
pow = find_output_node(model, pow.output[0])
|
||||
if not is_type(pow, "Pow"):
|
||||
continue
|
||||
cast_pow = find_input_node(model, pow.input[1])
|
||||
if not is_type(cast_pow, "Cast"):
|
||||
continue
|
||||
remove_nodes += [cast_pow]
|
||||
if not is_type(div, "Div") or not is_type(pow, "Pow"):
|
||||
continue
|
||||
|
||||
# check that pow ouput is ReduceMean
|
||||
reduce_mean2 = find_output_node(model, pow.output[0])
|
||||
if not is_type(reduce_mean2, "ReduceMean"):
|
||||
continue
|
||||
|
||||
# check that reduce_mean2 output is Add
|
||||
add = find_output_node(model, reduce_mean2.output[0])
|
||||
if not is_type(add, "Add"):
|
||||
continue
|
||||
|
||||
# check that add output is Sqrt
|
||||
sqrt = find_output_node(model, add.output[0])
|
||||
if not is_type(sqrt, "Sqrt"):
|
||||
continue
|
||||
|
||||
# check that sqrt output is div
|
||||
if div != find_output_node(model, sqrt.output[0]):
|
||||
continue
|
||||
|
||||
# check if div output is Mul
|
||||
optional_mul = find_output_node(model, div.output[0])
|
||||
if not is_type(optional_mul, "Mul"):
|
||||
optional_mul = None
|
||||
continue # default bias and weight not supported
|
||||
|
||||
# check if mul output is Add
|
||||
if optional_mul is not None:
|
||||
optional_add = find_output_node(model, optional_mul.output[0])
|
||||
else:
|
||||
optional_add = find_output_node(model, div.output[0])
|
||||
if not is_type(optional_add, "Add"):
|
||||
optional_add = None
|
||||
continue # default bias and weight not supported
|
||||
|
||||
|
||||
# add nodes to remove_nodes
|
||||
remove_nodes.extend([reduce_mean, sub, div, pow, reduce_mean2, add, sqrt])
|
||||
|
||||
# create LayerNorm node
|
||||
layer_norm_input = []
|
||||
layer_norm_output = []
|
||||
|
||||
layer_norm_input.append(reduce_mean.input[0])
|
||||
|
||||
if optional_mul is not None:
|
||||
remove_nodes.append(optional_mul)
|
||||
weight = optional_mul.input[1]
|
||||
layer_norm_input.append(weight)
|
||||
|
||||
if optional_add is not None:
|
||||
remove_nodes.append(optional_add)
|
||||
bias = optional_add.input[1]
|
||||
layer_norm_input.append(bias)
|
||||
|
||||
if optional_add is not None:
|
||||
layer_norm_output.append(optional_add.output[0])
|
||||
elif optional_mul is not None:
|
||||
layer_norm_output.append(optional_mul.output[0])
|
||||
else:
|
||||
layer_norm_output.append(div.output[0])
|
||||
|
||||
layer_norm_output.append('saved_mean_' + str(id))
|
||||
layer_norm_output.append('saved_inv_std_var_' + str(id))
|
||||
|
||||
epsilon_node = find_input_node(model, add.input[1])
|
||||
epsilon = epsilon_node.attribute[0].t.raw_data
|
||||
epsilon = struct.unpack('f', epsilon)[0]
|
||||
|
||||
layer_norm = helper.make_node("LayerNormalization",
|
||||
layer_norm_input,
|
||||
layer_norm_output,
|
||||
"LayerNormalization_" + str(id),
|
||||
None,
|
||||
axis = reduce_mean.attribute[0].ints[0],
|
||||
epsilon = epsilon)
|
||||
layer_norm_nodes.append(layer_norm)
|
||||
id += 1
|
||||
|
||||
# remove orphan constant nodes
|
||||
for constant in graph.node:
|
||||
if constant.op_type == "Constant" and constant not in remove_nodes:
|
||||
is_orphan = True
|
||||
for out_name in constant.output:
|
||||
out = find_output_node(model, out_name)
|
||||
if out not in remove_nodes:
|
||||
is_orphan = False
|
||||
if is_orphan:
|
||||
remove_nodes.append(constant)
|
||||
|
||||
all_nodes = []
|
||||
for node in graph.node:
|
||||
if node not in remove_nodes:
|
||||
all_nodes.append(node)
|
||||
|
||||
for node in layer_norm_nodes:
|
||||
all_nodes.append(node)
|
||||
|
||||
graph.ClearField("node")
|
||||
graph.node.extend(all_nodes)
|
||||
return model
|
||||
|
||||
# Fuse SoftmaxCrossEntropy
|
||||
|
||||
def fuse_softmaxNLL_to_softmaxCE(onnx_model):
|
||||
# Converting below subgraph
|
||||
#
|
||||
# (subgraph)
|
||||
# |
|
||||
# LogSoftmax (target) (optional weight)
|
||||
# | | |
|
||||
# nll_loss/NegativeLogLikelihoodLoss
|
||||
# |
|
||||
# output
|
||||
#
|
||||
# to the following
|
||||
#
|
||||
# (subgraph) (target) (optional weight)
|
||||
# | | _____|
|
||||
# | | |
|
||||
# SparseSoftmaxCrossEntropy
|
||||
# |
|
||||
# output
|
||||
nll_count = 0
|
||||
while True:
|
||||
nll_count = nll_count + 1
|
||||
nll_loss_node = None
|
||||
nll_loss_node_index = 0
|
||||
for nll_loss_node_index, node in enumerate(onnx_model.graph.node):
|
||||
if node.op_type == "nll_loss" or node.op_type == "NegativeLogLikelihoodLoss":
|
||||
nll_loss_node = node
|
||||
break
|
||||
|
||||
if nll_loss_node is None:
|
||||
break
|
||||
|
||||
softmax_node = None
|
||||
softmax_node_index = 0
|
||||
label_input_name = None
|
||||
weight_input_name = None
|
||||
for softmax_node_index, node in enumerate(onnx_model.graph.node):
|
||||
if node.op_type == "LogSoftmax":
|
||||
# has to be connected to nll_loss
|
||||
if len(nll_loss_node.input) > 2:
|
||||
weight_input_name = nll_loss_node.input[2]
|
||||
if node.output[0] == nll_loss_node.input[0]:
|
||||
softmax_node = node
|
||||
label_input_name = nll_loss_node.input[1]
|
||||
break
|
||||
elif node.output[0] == nll_loss_node.input[1]:
|
||||
softmax_node = node
|
||||
label_input_name = nll_loss_node.input[0]
|
||||
break
|
||||
else:
|
||||
if softmax_node is not None:
|
||||
break
|
||||
|
||||
if softmax_node is None:
|
||||
break
|
||||
|
||||
# delete nll_loss and LogSoftmax nodes in order
|
||||
if nll_loss_node_index < softmax_node_index:
|
||||
del onnx_model.graph.node[softmax_node_index]
|
||||
del onnx_model.graph.node[nll_loss_node_index]
|
||||
else:
|
||||
del onnx_model.graph.node[nll_loss_node_index]
|
||||
del onnx_model.graph.node[softmax_node_index]
|
||||
|
||||
probability_output_name = softmax_node.output[0]
|
||||
node = onnx_model.graph.node.add()
|
||||
inputs = [softmax_node.input[0], label_input_name, weight_input_name] if weight_input_name else [softmax_node.input[0], label_input_name]
|
||||
node.CopyFrom(onnx.helper.make_node("SparseSoftmaxCrossEntropy", inputs,
|
||||
[nll_loss_node.output[0], probability_output_name],
|
||||
"nll_loss_node_" + str(nll_count)))
|
||||
|
||||
return onnx_model
|
||||
|
|
@ -0,0 +1,214 @@
|
|||
import unittest
|
||||
import pytest
|
||||
import sys
|
||||
import os
|
||||
import copy
|
||||
from numpy.testing import assert_allclose, assert_array_equal
|
||||
|
||||
import onnx
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from orttraining_test_utils import map_optimizer_attributes
|
||||
from orttraining_test_transformers import BertModelTest, BertForPreTraining
|
||||
from orttraining_test_data_loader import create_ort_test_dataloader
|
||||
from orttraining_test_bert_postprocess import postprocess_model
|
||||
import onnxruntime
|
||||
|
||||
from onnxruntime.capi.ort_trainer import ORTTrainer, IODescription, ModelDescription, LossScaler, generate_sample
|
||||
|
||||
torch.manual_seed(1)
|
||||
onnxruntime.set_seed(1)
|
||||
|
||||
class Test_PostPasses(unittest.TestCase):
|
||||
def get_onnx_model(self, model, model_desc, inputs, device,
|
||||
_enable_internal_postprocess=True, _extra_postprocess=None):
|
||||
lr_desc = IODescription('Learning_Rate', [1,], torch.float32)
|
||||
model = ORTTrainer(model,
|
||||
None,
|
||||
model_desc,
|
||||
"LambOptimizer",
|
||||
map_optimizer_attributes,
|
||||
lr_desc,
|
||||
device,
|
||||
world_rank=0,
|
||||
world_size=1,
|
||||
_opset_version=12,
|
||||
_enable_internal_postprocess=_enable_internal_postprocess,
|
||||
_extra_postprocess=_extra_postprocess)
|
||||
|
||||
train_output = model.train_step(*inputs)
|
||||
return model.onnx_model_
|
||||
|
||||
def count_all_nodes(self, model):
|
||||
return len(model.graph.node)
|
||||
|
||||
def count_nodes(self, model, node_type):
|
||||
count = 0
|
||||
for node in model.graph.node:
|
||||
if node.op_type == node_type:
|
||||
count += 1
|
||||
return count
|
||||
|
||||
def find_nodes(self, model, node_type):
|
||||
nodes = []
|
||||
for node in model.graph.node:
|
||||
if node.op_type == node_type:
|
||||
nodes.append(node)
|
||||
return nodes
|
||||
|
||||
def get_name(self, name):
|
||||
if os.path.exists(name):
|
||||
return name
|
||||
rel = os.path.join("testdata", name)
|
||||
if os.path.exists(rel):
|
||||
return rel
|
||||
this = os.path.dirname(__file__)
|
||||
data = os.path.join(this, "..", "..", "..", "..", "onnxruntime", "test", "testdata")
|
||||
res = os.path.join(data, name)
|
||||
if os.path.exists(res):
|
||||
return res
|
||||
raise FileNotFoundError("Unable to find '{0}' or '{1}' or '{2}'".format(name, rel, res))
|
||||
|
||||
def test_layer_norm(self):
|
||||
class LayerNormNet(nn.Module):
|
||||
def __init__(self, target):
|
||||
super(LayerNormNet, self).__init__()
|
||||
self.ln_1 = nn.LayerNorm(10)
|
||||
self.loss = nn.CrossEntropyLoss()
|
||||
self.target = target
|
||||
|
||||
def forward(self, x):
|
||||
output1 = self.ln_1(x)
|
||||
loss = self.loss(output1, self.target)
|
||||
return loss, output1
|
||||
|
||||
device = torch.device("cpu")
|
||||
target = torch.ones(20, 10, 10, dtype=torch.int64).to(device)
|
||||
model = LayerNormNet(target)
|
||||
input = torch.randn(20, 5, 10, 10, dtype=torch.float32).to(device)
|
||||
|
||||
input_desc = IODescription('input', [], "float32")
|
||||
output0_desc = IODescription('output0', [], "float32")
|
||||
output1_desc = IODescription('output1', [20, 5, 10, 10], "float32")
|
||||
model_desc = ModelDescription([input_desc], [output0_desc, output1_desc])
|
||||
|
||||
learning_rate = torch.tensor([1.0000000e+00]).to(device)
|
||||
input_args=[input, learning_rate]
|
||||
|
||||
onnx_model = self.get_onnx_model(model, model_desc, input_args, device)
|
||||
|
||||
count_layer_norm = self.count_nodes(onnx_model, "LayerNormalization")
|
||||
count_nodes = self.count_all_nodes(onnx_model)
|
||||
|
||||
assert count_layer_norm == 1
|
||||
assert count_nodes == 3
|
||||
|
||||
def test_expand(self):
|
||||
class ExpandNet(nn.Module):
|
||||
def __init__(self, target):
|
||||
super(ExpandNet, self).__init__()
|
||||
self.loss = nn.CrossEntropyLoss()
|
||||
self.target = target
|
||||
self.linear = torch.nn.Linear(2, 2)
|
||||
|
||||
def forward(self, x, x1):
|
||||
output = x.expand_as(x1)
|
||||
output = self.linear(output)
|
||||
output = output + output
|
||||
loss = self.loss(output, self.target)
|
||||
return loss, output
|
||||
|
||||
device = torch.device("cpu")
|
||||
target = torch.ones(5, 5, 2, dtype=torch.int64).to(device)
|
||||
model = ExpandNet(target).to(device)
|
||||
|
||||
x = torch.randn(5, 3, 1, 2, dtype=torch.float32).to(device)
|
||||
x1 = torch.randn(5, 3, 5, 2, dtype=torch.float32).to(device)
|
||||
|
||||
input0_desc = IODescription('x', [5, 3, 1, 2], "float32")
|
||||
input1_desc = IODescription('x1', [5, 3, 5, 2], "float32")
|
||||
output0_desc = IODescription('output0', [], "float32")
|
||||
output1_desc = IODescription('output1', [5, 3, 5, 2], "float32")
|
||||
model_desc = ModelDescription([input0_desc, input1_desc], [output0_desc, output1_desc])
|
||||
|
||||
learning_rate = torch.tensor([1.0000000e+00]).to(device)
|
||||
input_args = [x, x1, learning_rate]
|
||||
|
||||
onnx_model = self.get_onnx_model(model, model_desc, input_args, device)
|
||||
|
||||
# check that expand output has shape
|
||||
expand_nodes = self.find_nodes(onnx_model, "Expand")
|
||||
assert len(expand_nodes) == 1
|
||||
|
||||
model_info = onnx_model.graph.value_info
|
||||
assert model_info[0].name == expand_nodes[0].output[0]
|
||||
assert model_info[0].type == onnx_model.graph.input[1].type
|
||||
|
||||
def test_bert(self):
|
||||
device = torch.device("cpu")
|
||||
|
||||
model_tester = BertModelTest.BertModelTester(self)
|
||||
config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels = model_tester.prepare_config_and_inputs()
|
||||
|
||||
model = BertForPreTraining(config=config)
|
||||
model.eval()
|
||||
|
||||
loss, prediction_scores, seq_relationship_score = model(input_ids,
|
||||
attention_mask=input_mask,
|
||||
token_type_ids=token_type_ids,
|
||||
masked_lm_labels=token_labels,
|
||||
next_sentence_label=sequence_labels)
|
||||
|
||||
model_desc = ModelDescription([model_tester.input_ids_desc,
|
||||
model_tester.attention_mask_desc,
|
||||
model_tester.token_type_ids_desc,
|
||||
model_tester.masked_lm_labels_desc,
|
||||
model_tester.next_sentence_label_desc],
|
||||
[model_tester.loss_desc,
|
||||
model_tester.prediction_scores_desc,
|
||||
model_tester.seq_relationship_scores_desc])
|
||||
|
||||
from collections import namedtuple
|
||||
MyArgs = namedtuple("MyArgs",
|
||||
"local_rank world_size max_steps learning_rate warmup_proportion batch_size seq_len")
|
||||
args = MyArgs(local_rank=0,
|
||||
world_size=1,
|
||||
max_steps=100,
|
||||
learning_rate=0.00001,
|
||||
warmup_proportion=0.01,
|
||||
batch_size=13,
|
||||
seq_len=7)
|
||||
|
||||
dataloader = create_ort_test_dataloader(model_desc.inputs_,
|
||||
args.batch_size,
|
||||
args.seq_len,
|
||||
device)
|
||||
learning_rate = torch.tensor(1.0e+0, dtype=torch.float32).to(device)
|
||||
for b in dataloader:
|
||||
batch = b
|
||||
break
|
||||
learning_rate = torch.tensor([1.00e+00]).to(device)
|
||||
inputs = batch + [learning_rate,]
|
||||
|
||||
onnx_model = self.get_onnx_model(model, model_desc, inputs, device, _extra_postprocess=postprocess_model)
|
||||
|
||||
self._bert_helper(onnx_model)
|
||||
|
||||
def _bert_helper(self, onnx_model):
|
||||
# count layer_norm
|
||||
count_layer_norm = self.count_nodes(onnx_model, "LayerNormalization")
|
||||
assert count_layer_norm == 12
|
||||
|
||||
# get expand node and check output shape
|
||||
expand_nodes = self.find_nodes(onnx_model, "Expand")
|
||||
assert len(expand_nodes) == 1
|
||||
|
||||
model_info = onnx_model.graph.value_info
|
||||
assert model_info[0].name == expand_nodes[0].output[0]
|
||||
assert model_info[0].type == onnx_model.graph.input[0].type
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main(module=__name__, buffer=True)
|
||||
|
|
@ -3,10 +3,3 @@ from orttraining_test_layer_norm_transform import layer_norm_transform
|
|||
|
||||
def postprocess_model(model):
|
||||
add_name(model)
|
||||
|
||||
# remove transpose node if its input is a 2d weight which only feeds to the node
|
||||
fix_transpose(model)
|
||||
|
||||
add_expand_shape(model)
|
||||
|
||||
layer_norm_transform(model)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ def warmup_linear(x, warmup=0.002):
|
|||
if x < warmup:
|
||||
return x/warmup
|
||||
return max((x - 1. )/ (warmup - 1.), 0.)
|
||||
|
||||
|
||||
def warmup_poly(x, warmup=0.002, degree=0.5):
|
||||
if x < warmup:
|
||||
return x/warmup
|
||||
|
|
@ -49,15 +49,16 @@ def map_optimizer_attributes(name):
|
|||
return {"alpha": 0.9, "beta": 0.999, "lambda": 0.01, "epsilon": 1e-6}
|
||||
|
||||
def run_test(model, model_desc, device, args, gradient_accumulation_steps, fp16,
|
||||
allreduce_post_accumulation, get_lr_this_step, use_internal_get_lr_this_step, loss_scaler, use_internal_loss_scaler,
|
||||
allreduce_post_accumulation, get_lr_this_step, use_internal_get_lr_this_step, loss_scaler, use_internal_loss_scaler,
|
||||
batch_args_option):
|
||||
dataloader = create_ort_test_dataloader(model_desc.inputs_, args.batch_size, args.seq_len, device)
|
||||
|
||||
model = ORTTrainer(model, None, model_desc, "LambOptimizer",
|
||||
map_optimizer_attributes=map_optimizer_attributes,
|
||||
learning_rate_description=IODescription('Learning_Rate', [1,], torch.float32),
|
||||
device=device, postprocess_model=postprocess_model,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps,
|
||||
device=device,
|
||||
_enable_internal_postprocess=True,
|
||||
gradient_accumulation_steps=gradient_accumulation_steps,
|
||||
# BertLAMB default initial settings: b1=0.9, b2=0.999, e=1e-6
|
||||
world_rank=args.local_rank, world_size=args.world_size,
|
||||
use_mixed_precision=fp16,
|
||||
|
|
|
|||
Loading…
Reference in a new issue