concrete.ml.onnx.onnx_model_manipulations

module concrete.ml.onnx.onnx_model_manipulations

Some code to manipulate models.

---

function simplify_onnx_model

simplify_onnx_model(onnx_model: ModelProto)

Simplify an ONNX model, removes unused Constant nodes and Identity nodes.

Args:

• onnx_model (onnx.ModelProto): the model to simplify.

---

function remove_unused_constant_nodes

remove_unused_constant_nodes(onnx_model: ModelProto)

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.

---

function remove_identity_nodes

remove_identity_nodes(onnx_model: ModelProto)

Remove identity nodes from a model.

Args:

• onnx_model (onnx.ModelProto): the model for which we want to remove Identity nodes.

---

function keep_following_outputs_discard_others

keep_following_outputs_discard_others(
    onnx_model: ModelProto,
    outputs_to_keep: Iterable[str]
)

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.

---

function remove_node_types

remove_node_types(onnx_model: ModelProto, op_types_to_remove: List[str])

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.

---

function clean_graph_after_node_name

clean_graph_after_node_name(
    onnx_model: ModelProto,
    node_name: str,
    fail_if_not_found: bool = True
)

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

---

function clean_graph_after_node_op_type

clean_graph_after_node_op_type(
    onnx_model: ModelProto,
    node_op_type: str,
    fail_if_not_found: bool = True
)

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