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

clean_graph_after_node(onnx_model: ModelProto, node_name: str)

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.