Coverage for mlprodict/onnx_tools/optim/onnx_optimisation.py: 100%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2@file
3@brief Optimisations of :epkg:`ONNX` graphs.
4"""
5from ._onnx_optimisation_common import _apply_optimisation_on_graph
6from .onnx_optimisation_identity import onnx_remove_node_identity
7from .onnx_optimisation_redundant import onnx_remove_node_redundant
8from .onnx_optimisation_unused import onnx_remove_node_unused
11def onnx_remove_node(onnx_model, recursive=True, debug_info=None, **options):
12 """
13 Removes as many nodes as possible without changing
14 the outcome. It applies @see fn onnx_remove_node_unused,
15 @see fn onnx_remove_node_identity,
16 and @see fn onnx_remove_node_redundant.
18 @param onnx_model onnx model
19 @param recursive looks into subgraphs
20 @param debug_info debug information (private)
21 @param options additional options
22 @return new onnx model
23 """
24 if debug_info is None:
25 debug_info = [str(type(onnx_model)).rsplit(
26 '.', maxsplit=1)[-1].strip("'>")]
27 else:
28 debug_info = (debug_info +
29 [str(type(onnx_model)).rsplit('.', maxsplit=1)[-1].strip("'>")])
31 if hasattr(onnx_model, 'graph'):
32 return _apply_optimisation_on_graph(
33 onnx_remove_node, onnx_model,
34 recursive=recursive, debug_info=debug_info,
35 **options)
37 graph = onnx_model
38 graph = onnx_remove_node_unused(
39 graph, recursive=recursive, debug_info=debug_info, **options)
40 graph = onnx_remove_node_identity(
41 graph, recursive=recursive, debug_info=debug_info, **options)
42 graph = onnx_remove_node_redundant(
43 graph, recursive=recursive, debug_info=debug_info, **options)
44 return graph