onnx_helpers ============ .. py:module:: onnx_helpers Classes ------- .. autoapisummary:: onnx_helpers.ModelSplit Functions --------- .. autoapisummary:: onnx_helpers.update_model_version onnx_helpers.remove_infer_shape onnx_helpers.load_model onnx_helpers.save_model onnx_helpers.run_model onnx_helpers.verify_models_equal onnx_helpers.verify_models_close onnx_helpers.extract_model onnx_helpers.split_model onnx_helpers.verify_split_models onnx_helpers.merge_model onnx_helpers.merge_split_model_with_shared_constant onnx_helpers.find_node onnx_helpers.find_node_output onnx_helpers.find_node_input onnx_helpers.make_node onnx_helpers.remove_node onnx_helpers.insert_node onnx_helpers.replace_node onnx_helpers.set_attribute_to_node onnx_helpers.remove_nodes_by_name_list onnx_helpers.insert_transpose_pair onnx_helpers.rewrite_matmul_as_conv onnx_helpers.rewrite_matmul_as_einsum onnx_helpers.insert_slices_after_node onnx_helpers.insert_concat onnx_helpers.insert_slices_concat_between_nodes onnx_helpers.rewrite_gemm_as_conv onnx_helpers.connect_nodes onnx_helpers.remove_output onnx_helpers.remove_outputs_by_name_list onnx_helpers.remove_inputs_by_name_list onnx_helpers.add_io onnx_helpers.add_input onnx_helpers.add_output onnx_helpers.get_io_shape onnx_helpers.update_io_shape onnx_helpers.change_node_output onnx_helpers.is_initializer onnx_helpers.find_initializer_value onnx_helpers.remove_initializers onnx_helpers.add_initializer onnx_helpers.remove_duplicated_initializer onnx_helpers.transpose_reshape_constant onnx_helpers.convert_initializer onnx_helpers.uniquify_initializer_name onnx_helpers.export_to_json onnx_helpers.update_model_node_names Module Contents --------------- .. py:class:: ModelSplit Configuration about how to split a model. :param split_name: Name of the extracted model. :param input_names: List of input names demarcating the beginning of the split. :param output_names: List of output names demarcating the end of the split. :param filename: File name of the extracted model to be saved. :param parent_split: ModelSplit for the parent if the parent is a split. Cannot be None when parent_path is None. :param parent_path: Path to the parent model to extract from. .. py:attribute:: split_name :type: str .. py:attribute:: input_names :type: List[str] .. py:attribute:: output_names :type: List[str] .. py:attribute:: filename :type: Optional[pathlib.Path] :value: None .. py:attribute:: parent_split :type: Optional[ModelSplit] :value: None .. py:attribute:: parent_path :type: Optional[pathlib.Path] :value: None .. py:method:: extract_from_parent_model() .. py:method:: run(input_data: List[numpy.ndarray]) -> List[numpy.ndarray] .. py:function:: update_model_version(model: onnx.ModelProto) -> onnx.ModelProto Update an ONNX model with IR and OPSET supported by Model SDK. :param model: Loaded ONNX model in onnx.ModelProto representation. :return: Updated ONNX model in onnx.ModelProto representation. .. py:function:: remove_infer_shape(model: onnx.ModelProto) Remove shape inference results in a model. :param model: Loaded model in onnx.ModelProto representation. .. py:function:: load_model(model_fname: str, load_only: bool = False) -> onnx.ModelProto Load a model and update its version information. :param model_fname: File name of the model to load from disk. :param load_only: Boolean flag, if set to False, to simplify the model after loading. :return: Loaded model in onnx.ModelProto representation. .. py:function:: save_model(model: onnx.ModelProto, model_fname: str, save_only: bool = False) Save a model to disk. :param model: Model to be saved. :param model_fname: File name to be used to save the model. :param save_only: Boolean flag, if set to False, to simplify and re-generate shape inference result before saving to disk. .. py:function:: run_model(model_name: str, input_names: List[str], input_data: List[numpy.ndarray]) -> List[numpy.ndarray] Run inference on a model. :param model_name: File name of the model saved on disk. :param input_names: List of input names expected by the model. :param input_data: List of input data for the model to run inference on. :return: List of outputs from inference result. .. py:function:: verify_models_equal(model_0: str, model_1: str, input_names: List[str], input_data: List[numpy.ndarray]) Run inference on two models and expect identical match numerically. :param model_0: File name of the first model. :param model_1: File name of the second model. :param input_names: List of input names expected by both models. :param input_data: List of input data to run inference on. .. py:function:: verify_models_close(model_0: str, model_1: str, input_names: List[str], input_data: List[numpy.ndarray], atol: float = 1e-06) Run inference on two models and expect close match numerically. :param model_0: File name of the first model. :param model_1: File name of the second model. :param input_names: List of input names expected by both models. :param input_data: List of input data to run inference on. :param atol: Float value as absolute tolerance in comparison. .. py:function:: extract_model(model: onnx.ModelProto, input_names: List[str], output_names: List[str]) -> onnx.ModelProto Extract from a model. :param model: The model to extract from. :param input_names: List of input names demarcating the beginning of the split. :param output_names: List of output names demarcating the end of the split. :return: Extracted model in onnx.ModelProto representation. .. py:function:: split_model(model_splits: List[ModelSplit]) Split a model. :param model_splits: List of ModelSplit to perform model extraction. .. py:function:: verify_split_models(whole_model: str, split_models: List[ModelSplit], input_names: List[str], input_data: List[numpy.ndarray]) Verify split models by comparing final inference outputs of the whole model and the cascaded splits. The order of the splits must be in the same execution order of the original whole model. :param whole_model: File name of the original whole model. :param split_models: List of ModelSplit extracted from the whole model. :param input_names: List of input names expected by the whole model. :param input_data: List of input data to run inference on. .. py:function:: merge_model(model_0: Union[str, pathlib.Path, onnx.ModelProto], model_1: Union[str, pathlib.Path, onnx.ModelProto], io_map: List[Tuple[str, str]]) -> onnx.ModelProto Merge two models. :param model_0: First model provided by file name or loaded as onnx.ModelProto. :param model_1: Second model provided by file name or loaded as onnx.ModelProto. :param io_map: List of tuples mapping output of the first model to input of the second model. :return: Merged model in onnx.ModelProto representation. .. py:function:: merge_split_model_with_shared_constant(base_model: Union[onnx.ModelProto, None], model_split: str, constant_prefix: str = None) -> onnx.ModelProto Merge two models which may have shared constants. When a model is split into multiple parts, some constants are duplicated in multiple parts. Later, when they are merged (possibly after modifications), there are conflicts because of the same constant names. A solution is to add prefix to constant names in the second model. :param model: Loaded model in onnx.ModelProto representation. :param model_split: File name of the second model to merge. :param constant_prefix: Prefix name to prepend to all constants in the second model. :return: Merged model in onnx.ModelProto representation. .. py:function:: find_node(model: onnx.ModelProto, node_name: str) -> onnx.NodeProto Find an operator node by name. :param model: Loaded model in onnx.ModelProto representation. :param node_name: Name of a node. :return: Operator node in onnx.NodeProto representation. .. py:function:: find_node_output(model: onnx.ModelProto, output_name: str) -> onnx.NodeProto Find an operator node whose output is the specified output_name. :param model: Loaded model in onnx.ModelProto representation. :param output_name: Name of the output. :return: Operator node in onnx.NodeProto representation. .. py:function:: find_node_input(model: onnx.ModelProto, input_name: str) -> Tuple[onnx.NodeProto, int] Find an operator node whose input is the specified input_name. :param model: Loaded model in onnx.ModelProto representation. :param input_name: Name of the input. :return: Tuple of Operator node and input index. .. py:function:: make_node(**kwargs: Any) -> onnx.NodeProto Make an operator node following ONNX specification. :param **kwargs: Operator dependent properties. :return: Newly constructed node in onnx.NodeProto representation. .. py:function:: remove_node(model: onnx.ModelProto, node: Union[str, onnx.NodeProto], remove_only: bool = False) Remove a node in a model. :param model: Loaded model in onnx.ModelProto representation. :param node: Node to be removed, provided as node name or onnx.NodeProto representation. :param remove_only: Boolean flag, if set to False, to reconnect surrounding nodes. .. py:function:: insert_node(model: onnx.ModelProto, node: onnx.NodeProto, new_node: onnx.NodeProto, insert_before: bool = False, insert_only: bool = False) Insert a node in a model. :param model: Loaded model in onnx.ModelProto representation. :param node: Reference node where to insert a new node. :param new_node: New node to be inserted. :param insert_before: Boolean flag, True to insert before or False to insert after the reference node. :param insert_only: Boolean flag, if False, to connect the new node to the reference node. .. py:function:: replace_node(model: onnx.ModelProto, node: onnx.NodeProto, new_nodes: List[onnx.NodeProto]) Replace a node in a model. :param model: Loaded model in onnx.ModelProto representation. :param node: Reference node to replace. :param new_nodes: List of nodes to replace with. .. py:function:: set_attribute_to_node(model: onnx.ModelProto, node_name: str, attr_name: str, val: Any) Set attribute of a node in a model. :param model: Loaded model in onnx.ModelProto representation. :param node_name: Name of the node. :param attr_name: Name of the attribute. :param val: Value of the attribute. .. py:function:: remove_nodes_by_name_list(model: onnx.ModelProto, name_list: List[str]) Remove multiple nodes by a list of names. :param model: Loaded model in onnx.ModelProto representation. :param name_list: List of names of the nodes to be removed. .. py:function:: insert_transpose_pair(model: onnx.ModelProto, node_name: str, perm_before: Tuple[int], perm_after: Tuple[int]) -> List[str] Insert a pair of transpose operators before and after a node. :param model: Loaded model in onnx.ModelProto representation. :param node_name: Name of the reference node. :param perm_before: Permutation before the node. :param perm_after: Permutation after the node. :return: List of names of the newly inserted transpose nodes. .. py:function:: rewrite_matmul_as_conv(model: onnx.ModelProto, name_list: List[str]) -> List[str] Rewrite a Matmul operator as a Conv operator. :param model: Loaded model in onnx.ModelProto representation. :param name_list: List of Matmul node names to be rewritten. :return: List of names of the new convolution nodes. .. py:function:: rewrite_matmul_as_einsum(model: onnx.ModelProto, eqn_list: Dict[str, str]) -> List[str] Rewrite a Matmul operator as an Einsum operator. :param model: Loaded model in onnx.ModelProto representation. :param eqn_list: Dictionary for rewrite with key being node name and value being equation string for Einsum. :return: List of names of the new Einsum nodes. .. py:function:: insert_slices_after_node(model: onnx.ModelProto, node_name: str, *, axis: int, nslices: int, slice_size: int) -> List[str] Insert Slice operators after a node. All Slice operators output the same length on the slicing axis. :param model: Loaded model in onnx.ModelProto representation. :param node_name: Name of the node after which to insert slices. :param axis: Axis for the Slice operators. :param nslices: Number of Slice operators. :param slice_size: Length for the Slice operators on the slicing axis. :return: List of names of the new Slice nodes. .. py:function:: insert_concat(model: onnx.ModelProto, input_nodes: List[str], at_node: str, *, axis: int) -> str Insert Concat operator after a group of nodes. :param model: Loaded model in onnx.ModelProto representation. :param input_nodes: List of names representing the group of nodes which are inputs to Concat. :param at_node: Name of node where insertion happens. :param axis: Axis on which to concatenate. :return: Name of newly inserted Concat node. .. py:function:: insert_slices_concat_between_nodes(model: onnx.ModelProto, after_node: str, before_node: str, *, slice_axis: int, nslices: int, slice_size: int, concat_axis: int) -> str Insert Slices + Concat between two nodes. :param model: Loaded model in onnx.ModelProto representation. :param after_node: Name of the node after which a group of Slice operators are inserted. :param before_node: Name of the node before which a Concat operator is inserted. :param slice_axis: Axis for the Slice operators. :param nslices: Number of Slice operators. :param slice_size: Length for the Slice operators on the slicing axis. :param concat_axis: Axis on which to concatenate. :return: Name of newly inserted Concat node. .. py:function:: rewrite_gemm_as_conv(model: onnx.ModelProto, after_node: str, at_node: str, *, w_r: int, w_s: int, w_c: int, w_k: int) -> str Rewrite a Gemm operator as a Conv operator between two nodes. :param model: Loaded model in onnx.ModelProto representation. :param after_node: Name of the node after which the Conv operator is inserted. :param at_node: Name of the node before which the Conv operator is inserted. :param w_r, w_s, w_c, w_k: Shape of the convolution weight as RSCK :return: Name of newly inserted Conv node. .. py:function:: connect_nodes(model: onnx.ModelProto, node_pair: List[str], out_idx: int, in_idx: int) Connect two nodes by feeding the output of one node to the input of next node. :param model: Loaded model in onnx.ModelProto representation. :param node_pair: List of two nodes in execution order. :param out_idx: The index of output of the first node. :param in_idx: The index of input of the second node. .. py:function:: remove_output(model: onnx.ModelProto, keep_list: List[str] = None) Run outputs of a model. :param model: Loaded model in onnx.ModelProto representation. :param keep_list: List of output names not to be removed. .. py:function:: remove_outputs_by_name_list(model: onnx.ModelProto, name_list: List[str]) Remove multiple outputs by a list of names. :param model: Loaded model in onnx.ModelProto representation. :param name_list: List of names of the outputs to be removed. .. py:function:: remove_inputs_by_name_list(model: onnx.ModelProto, name_list: List[str]) Remove multiple inputs by a list of names. :param model: Loaded model in onnx.ModelProto representation. :param name_list: List of names of the inputs to be removed. .. py:function:: add_io(model: onnx.ModelProto, io_name: str, io_shape: Tuple[int], io_dir: str) Add input or output to a model. :param model: Loaded model in onnx.ModelProto representation. :param io_name: Name to be added for an input/output. :param io_shape: Shape of the added input/output. :param io_dir: Designation of input or output. .. py:function:: add_input(model: onnx.ModelProto, input_name: str, input_shape: Tuple[int]) Add input to a model. :param model: Loaded model in onnx.ModelProto representation. :param input_name: Name of the input to be added. :param input_shape: Shape of the input to be added. .. py:function:: add_output(model: onnx.ModelProto, output_name, output_shape) Add output to a model. :param model: Loaded model in onnx.ModelProto representation. :param output_name: Name of the output to be added. :param output_shape: Shape of the output to be added. .. py:function:: get_io_shape(node: onnx.NodeProto) -> Tuple[int] Get the shape of a tensor. :param node: Node representing the tensor. :return: Shape of the tensor. .. py:function:: update_io_shape(model: onnx.ModelProto, io_name: str, new_shape: Tuple[int]) Update I/O shape in a model. :param model: Loaded model in onnx.ModelProto representation. :param io_name: Name of the I/O :param new_shape: Shape of the I/O .. py:function:: change_node_output(model: onnx.ModelProto, node_name: str, new_output_name: str, old_output_name: str = None) Change output of a node in a model. :param model: Loaded model in onnx.ModelProto representation. :param node_name: Name of a node. :param new_output_name: Name of the new output for the node. :param old_output_name: Name of previous output to replace. .. py:function:: is_initializer(model: onnx.ModelProto, name: str) -> bool Check if a name is an initializer. :param model: Loaded model in onnx.ModelProto representation. :param name: Name to check on. :return: True if provided name is an initializer. .. py:function:: find_initializer_value(model: onnx.ModelProto, initializer_name: str) -> numpy.ndarray Find value of an initializer in a model. :param model: Loaded model in onnx.ModelProto representation. :param initializer_name: Name of the initializer. :return: Value of the initializer. .. py:function:: remove_initializers(model: onnx.ModelProto, name_list: List[str]) Remove initializers from a list. :param model: Loaded model in onnx.ModelProto representation. :param name_list: List of initializer names. .. py:function:: add_initializer(model: onnx.ModelProto, initializer_name: str, initializer_value: numpy.ndarray) Add an initializer to a model. :param model: Loaded model in onnx.ModelProto representation. :param initializer_name: Name of the initializer. :param initializer_value: Value of the initializer. .. py:function:: remove_duplicated_initializer(model: onnx.ModelProto) Remove duplicated initializers in a model. :param model: Loaded model in onnx.ModelProto representation. .. py:function:: transpose_reshape_constant(model: onnx.ModelProto, param_name: str, perm: Tuple[int], new_shape: Tuple[int]) Perform transpose and/or reshape on an existing constant initializer. :param model: Loaded model in onnx.ModelProto representation. :param param_name: Name of the constant as an initializer. :param perm: Permutation to be performed. :param new_shape: New shape to be reshaped. .. py:function:: convert_initializer(model: onnx.ModelProto, init_name: str, convert_type: str, convert_func: Optional[Callable[[numpy.ndarray], numpy.ndarray]] = None) Convert an initializer in a model. :param model: Loaded model in onnx.ModelProto representation. :param init_name: Name of the initializer. :param convert_type: String describing format conversion. :param convert_func: Custom conversion function for "custom" convert_type. .. py:function:: uniquify_initializer_name(prefix_name: str, model: onnx.ModelProto) -> onnx.ModelProto Uniquify initializer names in a model by a prefix name. :param prefix_name: Name of the prefix to prepend to all initializer names. :param model: Loaded model in onnx.ModelProto representation. :return: Updated model in onnx.ModelProto representation. .. py:function:: export_to_json(model_path: str, json_file_path: str = 'model.json') Convert a model to a simpler json readable representation :param model: Loaded model in onnx.ModelProto representation. :param the path where to write the json data to .. py:function:: update_model_node_names(model_path: str, updated_model_path: str = 'updated_model.onnx') Update a model, with node names At times, when converting a model from pytorch to onnx the nodes names are not generated. This API helps in updating the nodes with generated names identified by the operator :param model_path: Input modelpath :param updated_model_path: Modified model save path