concrete.ml.onnx.onnx_model_manipulations
Some code to manipulate models.
simplify_onnx_model
Simplify an ONNX model, removes unused Constant nodes and Identity nodes.
Args:
onnx_model
(onnx.ModelProto): the model to simplify.
remove_unused_constant_nodes
Remove unused Constant nodes in the provided onnx model.
Args:
onnx_model
(onnx.ModelProto): the model for which we want to remove unused Constant nodes.
remove_identity_nodes
Remove identity nodes from a model.
Args:
onnx_model
(onnx.ModelProto): the model for which we want to remove Identity nodes.
keep_following_outputs_discard_others
Keep the outputs given in outputs_to_keep and remove the others from the model.
Args:
onnx_model
(onnx.ModelProto): the ONNX model to modify.
outputs_to_keep
(Iterable[str]): the outputs to keep by name.
remove_node_types
Remove unnecessary nodes from the ONNX graph.
Args:
onnx_model
(onnx.ModelProto): The ONNX model to modify.
op_types_to_remove
(List[str]): The node types to remove from the graph.
Raises:
ValueError
: Wrong replacement by an Identity node.
clean_graph_after_node_name
Clean the graph of the onnx model by removing nodes after the given node name.
Args:
onnx_model
(onnx.ModelProto): The onnx model.
node_name
(str): The node's name whose following nodes will be removed.
fail_if_not_found
(bool): If true, abort if the node name is not found
Raises:
ValueError
: if the node name is not found and if fail_if_not_found is set
clean_graph_after_node_op_type
Clean the graph of the onnx model by removing nodes after the given node type.
Args:
onnx_model
(onnx.ModelProto): The onnx model.
node_op_type
(str): The node's op_type whose following nodes will be removed.
fail_if_not_found
(bool): If true, abort if the node op_type is not found
Raises:
ValueError
: if the node op_type is not found and if fail_if_not_found is set