def build(self, *gradient_names):
"""Adds a clip grad norm sub graph to the onnx model."""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# add the necessary graph initializers
add_node_eps_name = graph_utils.generate_random_graph_name("add_eps")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(add_node_eps_name, onnx.TensorProto.FLOAT,
[1], [1e-6]))
max_norm_name = graph_utils.generate_random_graph_name("max_norm")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(max_norm_name, onnx.TensorProto.FLOAT, [1],
[self._max_norm]))
# perform gradient clipping
total_norm_name = self._reduce(*gradient_names)
adjusted_total_norm_name = self._add(total_norm_name,
add_node_eps_name)
clip_coef_name = self._clip(
self._div(max_norm_name, adjusted_total_norm_name))
return [
self._mul(grad_name, clip_coef_name)
for grad_name in gradient_names
]
def build(self, loss_input_name: str, target_name: typing.Optional[str] = "target"):
"""Adds a BCEWithLogitsLoss subgraph on top of an onnx model.
Creates a block that measures the binary cross entropy with logits between
loss_input_name and the target_name. This block combines Sigmoid layer
followed by a BCELoss.
Args:
loss_input_name: string input representing the loss input
target_name: string input representing the target
Returns:
Returns a string of the output name from the loss
"""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph initializers for pos_weight, weight, and the sub operands ([1])
pos_weight_name = graph_utils.generate_random_graph_name("bceloss.pos_weight")
if self._pos_weight is not None:
onnx_model.graph.initializer.append(onnx.numpy_helper.from_array(self._pos_weight, pos_weight_name))
weight_name = graph_utils.generate_random_graph_name("bceloss.weight")
if self._weight is not None:
onnx_model.graph.initializer.append(onnx.numpy_helper.from_array(self._weight, weight_name))
sub_ones_operand_name1 = graph_utils.generate_random_graph_name("bceloss.sub_ones")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(sub_ones_operand_name1, onnx.TensorProto.FLOAT, [1], [1.0])
)
sub_ones_operand_name2 = graph_utils.generate_random_graph_name("bceloss.sub_ones")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(sub_ones_operand_name2, onnx.TensorProto.FLOAT, [1], [1.0])
)
# create a new graph input. this is the target input needed to compare
# the graph output against to calculate loss.
target_input = copy.deepcopy(graph_utils.get_output_from_output_name(onnx_model, loss_input_name))
target_input.name = target_name
onnx_model.graph.input.append(target_input)
# create the bceloss
sigmoid_output = self._sigmoid(loss_input_name)
add_operand1 = self._mul(self._log(sigmoid_output), target_name)
if self._pos_weight is not None:
add_operand1 = self._mul(add_operand1, pos_weight_name)
add_operand2 = self._mul(
self._log(self._sub(sub_ones_operand_name1, sigmoid_output)),
self._sub(sub_ones_operand_name2, target_name),
)
loss_output = self._neg(self._add(add_operand1, add_operand2))
if self._weight is not None:
loss_output = self._mul(weight_name, loss_output)
return self._reduce(loss_output)
def build(self, scores_input_name: str, labels_name: str = "labels"):
"""Adds a CrossEntropyLoss subgraph on top of an onnx model.
Creates a block that measures the softmax cross entropy between
scores_input_name and the labels_name.
Args:
loss_input_name: string input representing the loss input
labels_name: string input representing the labels
Returns:
Returns a string of the output name from the loss
"""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
weight_name = graph_utils.generate_random_graph_name("celoss.weight")
if self._weight is not None:
onnx_model.graph.initializer.append(onnx.numpy_helper.from_array(self._weight, weight_name))
# create a new graph input. this is the labels input needed to compare
# the graph output against to calculate loss.
labels_input = copy.deepcopy(graph_utils.get_output_from_output_name(onnx_model, scores_input_name))
labels_input.name = labels_name
labels_input.type.tensor_type.elem_type = onnx.TensorProto.INT32
# if the predictions are (num_examples x num_classes)
# labels should be (num_examples x 1)
del labels_input.type.tensor_type.shape.dim[1]
onnx_model.graph.input.append(labels_input)
# create a new graph node for the loss
loss_node_input_names = [scores_input_name, labels_name]
if self._weight:
loss_node_input_names.append(weight_name)
loss_node_output_name = graph_utils.generate_random_graph_name("loss")
loss_node_output_names = [
loss_node_output_name,
graph_utils.generate_random_graph_name("log_prob"),
]
loss_node = onnx.helper.make_node(
"SoftmaxCrossEntropyLoss",
loss_node_input_names,
loss_node_output_names,
reduction=self._reduction,
ignore_index=self._ignore_index,
name=graph_utils.generate_random_graph_name("SoftmaxCrossEntropyLoss"),
)
onnx_model.graph.node.append(loss_node)
return loss_node_output_name
def build( # pylint: disable=too-many-arguments
self,
learning_rate_name: str,
step_name: str,
parameter_sequence_name: str,
gradient_sequence_name: str,
first_order_moment_sequence_name: str,
second_order_moment_sequence_name: str,
):
"""Adds the AdamWOptimizer node to the model."""
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# define the node attributes
node_attributes = {
"alpha": self._betas[0], # beta1
"beta": self._betas[1], # beta2
"epsilon": self._eps, # epsilon
"weight_decay": self._weight_decay, # weight decay
"correct_bias":
1 if self._bias_correction else 0, # bias_correction
"adam_mode": 1, # adam mode (1 for hf/transformers/AdamW)
}
# add the adamw node to the onnx model
adamw_input_names = [
learning_rate_name, # learning rate
step_name, # training step
parameter_sequence_name, # param to be updated
gradient_sequence_name, # gradient of the param to be used for update
first_order_moment_sequence_name, # first order moment for this param
second_order_moment_sequence_name, # second order moment for this param
]
adamw_output_name = graph_utils.generate_random_graph_name(
"adamw.updated_flag")
adamw_output_names = [adamw_output_name]
adamw_node = onnx.helper.make_node(
"AdamWOptimizer",
adamw_input_names,
adamw_output_names,
name=graph_utils.generate_random_graph_name("AdamWOptimizer"),
domain="com.microsoft",
**node_attributes,
)
onnx_model.graph.node.append(adamw_node)
return adamw_output_name
def build(self):
# create the graph initializer for the exponent
initializer_name = graph_utils.generate_random_graph_name("initializer")
accessor.global_accessor.model.graph.initializer.append(
onnx.helper.make_tensor(initializer_name, onnx.TensorProto.FLOAT, [1], [self._value])
)
return initializer_name
def build(self, *sequence_input_names):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph node for this sequence construct node
sc_node_input_names = list(sequence_input_names)
sc_node_output_name = graph_utils.generate_random_graph_name("sequenceconstruct_output")
sc_node_output_names = [sc_node_output_name]
sc_node = onnx.helper.make_node(
"SequenceConstruct",
sc_node_input_names,
sc_node_output_names,
graph_utils.generate_random_graph_name("SequenceConstruct"),
)
onnx_model.graph.node.append(sc_node)
return sc_node_output_name
def build(self, input_name1, input_name2):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# Assert that the op name is not empty
if not self._op_name:
raise RuntimeError("Unknown op name. Please override _op_name")
# create the graph node for sub
node_input_names = [input_name1, input_name2]
node_output_name = graph_utils.generate_random_graph_name(f"{self._op_name.lower()}_output")
node_output_names = [node_output_name]
node = onnx.helper.make_node(
self._op_name,
node_input_names,
node_output_names,
name=graph_utils.generate_random_graph_name(self._op_name),
)
onnx_model.graph.node.append(node)
return node_output_name
def build(self, *reduce_node_input_names):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph node for this reducealll2 node
reduce_node_input_names = list(reduce_node_input_names)
reduce_node_output_name = graph_utils.generate_random_graph_name("reducealll2_output")
reduce_node_output_names = [reduce_node_output_name]
reduce_node = onnx.helper.make_node(
"ReduceAllL2",
reduce_node_input_names,
reduce_node_output_names,
graph_utils.generate_random_graph_name("ReduceAllL2"),
domain="com.microsoft",
)
onnx_model.graph.node.append(reduce_node)
# TODO: register shape inference with onnx
onnx_model.graph.value_info.append(
onnx.helper.make_tensor_value_info(reduce_node_output_name, onnx.TensorProto.FLOAT, [1])
)
return reduce_node_output_name
def build(self, pow_input_name):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph initializer for the exponent
pow_node_exponent_name = graph_utils.generate_random_graph_name("pow_exponent")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(pow_node_exponent_name, onnx.TensorProto.FLOAT, [1], [self._exponent])
)
# create the graph node for pow
pow_node_input_names = [pow_input_name, pow_node_exponent_name]
pow_node_output_name = graph_utils.generate_random_graph_name("pow_output")
pow_node_output_names = [pow_node_output_name]
pow_node = onnx.helper.make_node(
"Pow",
pow_node_input_names,
pow_node_output_names,
name=graph_utils.generate_random_graph_name("Pow"),
)
onnx_model.graph.node.append(pow_node)
return pow_node_output_name
def build(self, clip_input_name):
# get the model to manipulate
onnx_model = accessor.global_accessor.model
# create the graph initializer for the clip min
clip_node_min_name = ""
if self._min is not None:
clip_node_min_name = graph_utils.generate_random_graph_name("clip_min")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(clip_node_min_name, onnx.TensorProto.FLOAT, [1], [self._min])
)
# create the graph initializer for the clip max
clip_node_max_name = ""
if self._max is not None:
clip_node_max_name = graph_utils.generate_random_graph_name("clip_max")
onnx_model.graph.initializer.append(
onnx.helper.make_tensor(clip_node_max_name, onnx.TensorProto.FLOAT, [1], [self._max])
)
# create the graph node for this clip node
clip_node_input_names = [
clip_input_name,
clip_node_min_name,
clip_node_max_name,
]
clip_node_output_name = graph_utils.generate_random_graph_name("clip_output")
clip_node_output_names = [clip_node_output_name]
clip_node = onnx.helper.make_node(
"Clip",
clip_node_input_names,
clip_node_output_names,
graph_utils.generate_random_graph_name("Clip"),
)
onnx_model.graph.node.append(clip_node)
return clip_node_output_name