Skip to content

Commit

Permalink
add node to container instead of changing output node of container
Browse files Browse the repository at this point in the history
  • Loading branch information
janjagusch committed Jul 2, 2021
1 parent 93e95db commit c1fbebc
Showing 1 changed file with 8 additions and 16 deletions.
24 changes: 8 additions & 16 deletions onnxmltools/convert/lightgbm/operator_converters/LightGbm.py
Original file line number Diff line number Diff line change
Expand Up @@ -399,26 +399,18 @@ def convert_lightgbm(scope, operator, container):

apply_div(scope, [output_name, denominator_name],
operator.output_full_names, container, broadcast=1)
elif post_transform:
container.add_node(
post_transform,
output_name,
operator.output_full_names,
name=scope.get_unique_operator_name(
post_transform),
)
else:
container.add_node('Identity', output_name,
operator.output_full_names,
name=scope.get_unique_operator_name('Identity'))
if post_transform:
_add_post_transform_node(container, post_transform)


def _add_post_transform_node(container: ModelComponentContainer, op_type: str):
"""
Add a post transform node to a ModelComponentContainer.
Useful for post transform functions that are not supported by the ONNX spec yet (e.g. 'Exp').
"""
assert len(container.outputs) == 1, "Adding a post transform node is only possible for models with 1 output."
original_output_name = container.outputs[0].name
new_output_name = f"{op_type.lower()}_{original_output_name}"
post_transform_node = onnx.helper.make_node(op_type, inputs=[original_output_name], outputs=[new_output_name])
container.nodes.append(post_transform_node)
container.outputs[0].name = new_output_name


def modify_tree_for_rule_in_set(gbm, use_float=False):
Expand Down

0 comments on commit c1fbebc

Please sign in to comment.