ml-sdk/docs/references/api/concrete.ml.onnx.onnx_model_manipulations.md

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_at_node_op_type

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

Remove the first node matching node_op_type and its following nodes from the ONNX graph.

Args:

• onnx_model (onnx.ModelProto): The onnx model.
• node_op_type (str): The node's op_type whose following nodes as well as itself will be removed.
• fail_if_not_found (bool): If true, raise an error if no node matched the given op_type.

Raises:

• ValueError: If no node matched the given op_type and fail_if_not_found is set to True

---

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
)

Remove the nodes following first node matching node_op_type from the ONNX graph.

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, raise an error if no node matched the given op_type.

Raises:

• ValueError: If no node matched the given op_type and fail_if_not_found is set to True

---

function convert_first_gather_to_matmul

convert_first_gather_to_matmul(
    onnx_model: ModelProto
) → Tuple[Union[ModelProto, NoneType], ModelProto]

Convert the first Gather node to a matrix multiplication node.

In FHE, Gather is a costly operation since it can involve many PBS. When it appears first in the onnx model, we can remove it and replace it by a matrix multiplication node by converting the indices to a one-hot encoding.

Args:

• onnx_model (onnx.ModelProto): The onnx model.

Returns:

• Tuple[Optional[onnx.ModelProto], onnx.ModelProto]: The pre-processing model and the modified onnx model.