Coverage for mlprodict/onnx_tools/onnx_export.py: 99%
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 Exports an ONNX graph in a way it can we created again
4with a python script. It relies on :epkg:`jinja2` and :epkg:`autopep8`.
6.. versionadded:: 0.7
7"""
8from textwrap import indent
9import numpy
10import onnx
11from onnx import numpy_helper
12from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
13from .onnx2py_helper import (
14 _var_as_dict, guess_proto_dtype, guess_proto_dtype_name)
15from .onnx_export_templates import (
16 get_onnx_template, get_tf2onnx_template, get_numpy_template,
17 get_xop_template)
18from .exports.numpy_helper import make_numpy_code
19from .exports.tf2onnx_helper import make_tf2onnx_code
22def select_attribute(ens, att, sort=False, unique=False, skip=None):
23 """
24 Returns the list of the same attribute.
25 `[el.att for el in ens]`.
27 :param ens: list
28 :param att: attribute name
29 :param sort: sort the array
30 :param unique: returns the unique values
31 :param skip: to skip some names
32 :return: something like `[el.att for el in ens]`
33 """
34 if len(ens) == 0:
35 return []
36 if isinstance(ens[0], dict):
37 atts = [el[att] for el in ens]
38 else:
39 atts = [getattr(el, att) for el in ens]
40 if unique:
41 atts = list(set(atts))
42 if sort:
43 atts.sort()
44 if skip is None:
45 return atts
46 return [a for a in atts if a not in skip]
49def _nodes(graph, rename_name, used, output_names, use_onnx_tensor,
50 templates, verbose, opset, rename, autopep_options, name,
51 subgraphs, unique_operators):
52 from ..npy.xop import loadop
53 nodes = []
54 for node in graph.node:
55 if node.domain in ('', 'ai.onnx.ml'):
56 clname = loadop((node.domain, node.op_type))
57 unique_operators.add(
58 (node.domain, node.op_type, clname.__name__))
59 for index_input, i_raw_name in enumerate(node.input):
60 if len(i_raw_name) == 0:
61 # This means the input is optional.
62 if any(map(lambda s: len(s) > 0, node.input[index_input:])):
63 raise NotImplementedError(
64 "Input cannot be placed after an unused optional input "
65 "in node %r." % (node, ))
66 break
67 i = rename_name(i_raw_name)
68 if i not in used:
69 used[i] = []
70 used[i].append(node)
71 attributes = []
72 for at in node.attribute:
73 temp = _var_as_dict(at)
74 value = temp['value']
75 if node.op_type == 'Scan' and at.name == 'body':
76 fname = "_create_" + node.name + "_body"
77 body = export_template(
78 value, templates, opset=opset, verbose=verbose,
79 name=name, rename=rename,
80 use_onnx_tensor=use_onnx_tensor,
81 autopep_options=autopep_options,
82 function_name=fname)
83 subgraphs.append((body, node.name + "_body"))
84 attributes.append((at.name, fname + "()"))
85 continue
86 if node.op_type in {'Loop', 'If'}:
87 raise NotImplementedError(
88 "Subgraphs are not yet implemented (operator=%r)."
89 "" % node.op_type)
90 if use_onnx_tensor:
91 if node.op_type == 'Cast' and at.name == 'to':
92 attributes.append(
93 (at.name, guess_proto_dtype_name(int(value))))
94 continue
95 if isinstance(value, str):
96 attributes.append((at.name, "%r" % value))
97 else:
98 if isinstance(value, numpy.ndarray):
99 if use_onnx_tensor and at.name == 'value':
100 onnx_dtype = guess_proto_dtype_name(
101 guess_proto_dtype(value.dtype))
102 value = (
103 'make_tensor("value", %s, dims=%r, vals=%r)'
104 '' % (onnx_dtype, list(value.shape),
105 value.tolist()))
106 attributes.append((at.name, value))
107 else:
108 attributes.append((at.name, repr(value.tolist())))
109 else:
110 attributes.append((at.name, repr(value)))
112 attributes_str = ", ".join("%s=%s" % (k, v) for k, v in attributes)
113 d = dict(name=node.name, op_type=node.op_type,
114 domain=node.domain,
115 inputs=[rename_name(n) for n in node.input if len(n) > 0],
116 outputs=[rename_name(n) for n in node.output],
117 output_names=[rename_name(n) for n in node.output
118 if n in output_names],
119 attributes=attributes, attributes_str=attributes_str)
120 nodes.append(d)
121 return nodes
124def export_template(model_onnx, templates, opset=None, # pylint: disable=R0914
125 verbose=True, name=None,
126 rename=False, use_onnx_tensor=False,
127 autopep_options=None, function_name='create_model'):
128 """
129 Exports an ONNX model to the onnx syntax.
131 :param model_onnx: string or ONNX graph
132 :param templates: exporting templates
133 :param opset: opset to export to
134 (None to select the one from the graph)
135 :param verbose: insert prints
136 :param name: to overwrite onnx name
137 :param rename: rename the names to get shorter names
138 :param use_onnx_tensor: when an attribute is an array
139 and its name is `'value'`, it converts that array into an
140 ONNX tensor to avoid type mismatch, (operator *ConstantOfShape*, ...)
141 :param autopep_options: :epkg:`autopep8` options
142 :param function_name: main function name in the code
143 :return: python code
144 """
145 # delayed import to avoid raising an exception if not installed.
146 import autopep8
148 def number2name(n):
149 n += 1
150 seq = []
151 while n >= 1:
152 r = n % 26
153 seq.append(r)
154 n = (n - r) // 26
155 return "".join(chr(65 + i) for i in reversed(seq))
157 def rename_name(name):
158 if len(name) == 0:
159 raise ValueError( # pragma: no cover
160 "name is empty.")
161 if name in dict_names:
162 return dict_names[name]
163 if rename:
164 i = 0
165 new_name = number2name(i)
166 while new_name in dict_names:
167 i += 1
168 new_name = number2name(i)
169 if len(new_name) == 0:
170 raise ValueError( # pragma: no cover
171 "Unable to rename name=%r i=%d." % (name, i))
172 dict_names[name] = new_name
173 dict_names[new_name] = new_name
174 return new_name
175 return name
177 # containers
178 context = {}
179 used = {}
181 # opset
182 if hasattr(model_onnx, 'opset_import'):
183 opsets = {}
184 for oimp in model_onnx.opset_import:
185 if oimp.domain == '' and opset is None:
186 opsets[oimp.domain] = oimp.version
187 opset = oimp.version
188 else:
189 opsets[oimp.domain] = opset
190 context['opsets'] = opsets
191 context['target_opset'] = opset
193 if hasattr(model_onnx, 'graph'):
194 graph = model_onnx.graph
195 else:
196 graph = model_onnx
197 dict_names = {}
198 if rename:
199 for o in graph.input:
200 dict_names[o.name] = o.name
201 for o in graph.output:
202 dict_names[o.name] = o.name
204 # inits
205 unique_operators = set()
206 initializers = []
207 for init in graph.initializer:
208 init_name = rename_name(init.name)
209 value = numpy_helper.to_array(init)
210 initializers.append((init_name, value))
211 context['initializers'] = initializers
212 context['initializers_dict'] = {k: v for k, v in initializers}
214 # functions
215 functions = []
216 fct_dict = {}
217 if hasattr(model_onnx, 'functions'):
218 from ..npy.xop import OnnxOperatorFunction
219 for fct in model_onnx.functions:
220 used = {}
221 functions.append(
222 (fct.domain, fct.name,
223 {'proto': fct,
224 'nodes': _nodes(fct, rename_name, used, fct.output,
225 use_onnx_tensor, templates, verbose,
226 opset, rename, autopep_options,
227 fct.name, [], unique_operators)}))
228 if fct.name in fct_dict:
229 fct_dict[fct.name].append(fct)
230 else:
231 fct_dict[fct.name] = [fct]
232 context['OnnxOperatorFunction'] = OnnxOperatorFunction
233 context['functions'] = functions
234 context['functions_dict'] = fct_dict
236 # inputs
237 inputs = []
238 for inp in graph.input:
239 t = inp.type.tensor_type
240 dims = []
241 for d in t.shape.dim:
242 dd = d.dim_value
243 if dd == 0:
244 dd = None
245 dims.append(dd)
246 if len(dims) == 0:
247 dims = None
248 if 'dim_value' in str(dims):
249 raise RuntimeError( # pragma: no cover
250 "Unexpected issue in %r - %r." % (dims, t))
251 inputs.append((inp.name, t.elem_type, dims))
252 context['inputs'] = inputs
254 # outputs
255 outputs = []
256 for inp in graph.output:
257 t = inp.type.tensor_type
258 dims = []
259 for d in t.shape.dim:
260 dd = d.dim_value
261 if dd == 0:
262 dd = None
263 dims.append(dd)
264 if len(dims) == 0:
265 dims = None
266 outputs.append((inp.name, t.elem_type, dims))
267 context['outputs'] = outputs
269 # node
270 output_names = set(o.name for o in graph.output)
271 subgraphs = []
272 context['nodes'] = _nodes(
273 graph, rename_name, used, output_names, use_onnx_tensor,
274 templates, verbose, opset, rename, autopep_options, name,
275 subgraphs, unique_operators)
277 # graph
278 context['name'] = name or graph.name
279 context['name'] = context['name'].replace("(", "_").replace(")", "")
280 context['function_name'] = function_name
281 context['indent'] = indent
282 if hasattr(model_onnx, 'graph'):
283 context['ir_version'] = model_onnx.ir_version
284 context['producer_name'] = model_onnx.producer_name
285 context['domain'] = model_onnx.domain
286 context['model_version'] = model_onnx.model_version
287 context['doc_string'] = model_onnx.doc_string
288 context['metadata'] = {
289 p.key: p.value for p in model_onnx.metadata_props}
290 else:
291 # subgraph
292 context['ir_version'] = None
293 context['producer_name'] = None
294 context['domain'] = None
295 context['model_version'] = None
296 context['doc_string'] = ""
297 context['metadata'] = {}
299 # common context
300 context['unique_operators'] = [dict(domain=o[0], name=o[1], classname=o[2])
301 for o in sorted(unique_operators)]
302 context['skip_inits'] = {}
303 context['subgraphs'] = subgraphs
305 mark_inits = {}
307 # First rendering to detect any unused or replaced initializer.
308 from jinja2 import Template # delayed import
309 template = Template(templates)
310 final = template.render(
311 enumerate=enumerate, sorted=sorted, len=len,
312 select_attribute=select_attribute, repr=repr,
313 TENSOR_TYPE_TO_NP_TYPE=TENSOR_TYPE_TO_NP_TYPE,
314 make_numpy_code=lambda *args, **kwargs: make_numpy_code(
315 *args, context=context, used=used, mark_inits=mark_inits,
316 **kwargs),
317 make_tf2onnx_code=lambda *args, **kwargs: make_tf2onnx_code(
318 *args, context=context, used=used, mark_inits=mark_inits,
319 **kwargs),
320 verbose=verbose, **context)
322 skip_inits = set()
323 for k, v in mark_inits.items():
324 if len(v) == len(used[k]):
325 # One initializers was removed.
326 skip_inits.add(k)
328 if len(skip_inits) > 0:
329 # Second rendering if needed when an initializer was replaced
330 # or removed.
331 context['skip_inits'] = skip_inits
332 # Again with skip_inits.
333 final = template.render(
334 enumerate=enumerate, sorted=sorted, len=len,
335 make_numpy_code=lambda *args, **kwargs: make_numpy_code(
336 *args, context=context, used=used, mark_inits=mark_inits,
337 **kwargs),
338 make_tf2onnx_code=lambda *args, **kwargs: make_tf2onnx_code(
339 *args, context=context, used=used, mark_inits=mark_inits,
340 **kwargs),
341 verbose=verbose, **context)
343 final += "\n"
344 if not verbose:
345 rows = final.split("\n")
346 final = "\n".join(_ for _ in rows if not _.endswith("# verbose"))
347 return autopep8.fix_code(final, options=autopep_options)
350def export2onnx(model_onnx, opset=None, verbose=True, name=None, rename=False,
351 autopep_options=None):
352 """
353 Exports an ONNX model to the :epkg:`onnx` syntax.
355 :param model_onnx: string or ONNX graph
356 :param opset: opset to export to
357 (None to select the one from the graph)
358 :param verbose: inserts prints
359 :param name: to overwrite onnx name
360 :param rename: rename the names to get shorter names
361 :param autopep_options: :epkg:`autopep8` options
362 :return: python code
364 The following example shows what a python code creating a graph
365 implementing the KMeans would look like.
367 .. runpython::
368 :showcode:
369 :process:
371 import numpy
372 from sklearn.cluster import KMeans
373 from mlprodict.onnx_conv import to_onnx
374 from mlprodict.onnx_tools.onnx_export import export2onnx
376 X = numpy.arange(20).reshape(10, 2).astype(numpy.float32)
377 tr = KMeans(n_clusters=2)
378 tr.fit(X)
380 onx = to_onnx(tr, X, target_opset=14)
381 code = export2onnx(onx)
383 print(code)
384 """
385 if isinstance(model_onnx, str):
386 model_onnx = onnx.load(model_onnx)
388 code = export_template(model_onnx, templates=get_onnx_template(),
389 opset=opset, verbose=verbose, name=name,
390 rename=rename, use_onnx_tensor=True,
391 autopep_options=autopep_options)
392 return code
395def export2tf2onnx(model_onnx, opset=None, verbose=True, name=None,
396 rename=False, autopep_options=None):
397 """
398 Exports an ONNX model to the :epkg:`tensorflow-onnx` syntax.
400 :param model_onnx: string or ONNX graph
401 :param opset: opset to export to
402 (None to select the one from the graph)
403 :param verbose: inserts prints
404 :param name: to overwrite onnx name
405 :param rename: rename the names to get shorter names
406 :param autopep_options: :epkg:`autopep8` options
407 :return: python code
409 .. runpython::
410 :showcode:
411 :process:
413 import numpy
414 from sklearn.cluster import KMeans
415 from mlprodict.onnx_conv import to_onnx
416 from mlprodict.onnx_tools.onnx_export import export2tf2onnx
418 X = numpy.arange(20).reshape(10, 2).astype(numpy.float32)
419 tr = KMeans(n_clusters=2)
420 tr.fit(X)
422 onx = to_onnx(tr, X, target_opset=14)
423 code = export2tf2onnx(onx)
425 print(code)
426 """
427 if isinstance(model_onnx, str):
428 model_onnx = onnx.load(model_onnx)
430 code = export_template(model_onnx, templates=get_tf2onnx_template(),
431 opset=opset, verbose=verbose, name=name,
432 rename=rename, use_onnx_tensor=True,
433 autopep_options=autopep_options)
434 code = code.replace("], ]", "]]")
435 return code
438def export2numpy(model_onnx, opset=None, verbose=True, name=None,
439 rename=False, autopep_options=None):
440 """
441 Exports an ONNX model to the :epkg:`numpy` syntax.
442 The exports does not work with all operators.
444 :param model_onnx: string or ONNX graph
445 :param opset: opset to export to
446 (None to select the one from the graph)
447 :param verbose: inserts prints
448 :param name: to overwrite onnx name
449 :param rename: rename the names to get shorter names
450 :param autopep_options: :epkg:`autopep8` options
451 :return: python code
453 .. runpython::
454 :showcode:
455 :process:
457 import numpy
458 from sklearn.cluster import KMeans
459 from mlprodict.onnx_conv import to_onnx
460 from mlprodict.onnx_tools.onnx_export import export2numpy
462 X = numpy.arange(20).reshape(10, 2).astype(numpy.float32)
463 tr = KMeans(n_clusters=2)
464 tr.fit(X)
466 onx = to_onnx(tr, X, target_opset=14)
467 code = export2numpy(onx)
469 print(code)
471 This can be applied to the decomposition of an einsum
472 equation into simple matrix operations.
474 .. runpython::
475 :showcode:
476 :process:
478 import numpy
479 from mlprodict.testing.einsum import decompose_einsum_equation
480 from mlprodict.onnx_tools.onnx_export import export2numpy
482 x1 = numpy.arange(8).reshape(2, 2, 2).astype(numpy.float32)
483 x2 = numpy.arange(4).reshape(2, 2).astype(numpy.float32)
484 r = numpy.einsum("bac,cd->ad", x1, x2)
486 seq_clean = decompose_einsum_equation(
487 "bac,cd->ad", strategy='numpy', clean=True)
488 onx = seq_clean.to_onnx("Y", "X1", "X2", dtype=numpy.float32)
489 code = export2numpy(onx, name="einsum")
490 print(code)
491 """
492 if isinstance(model_onnx, str):
493 model_onnx = onnx.load(model_onnx)
495 code = export_template(model_onnx, templates=get_numpy_template(),
496 opset=opset, verbose=verbose, name=name,
497 rename=rename, autopep_options=autopep_options)
498 for i in range(-6, 6):
499 code = code.replace("axis=tuple([%d])" % i, "axis=%d" % i)
500 code = code.replace("tuple([%d])" % i, "(%d, )" % i)
501 return code
504def export2xop(model_onnx, opset=None, verbose=True, name=None, rename=False,
505 autopep_options=None):
506 """
507 Exports an ONNX model to the :epkg:`onnx` syntax.
509 :param model_onnx: string or ONNX graph
510 :param opset: opset to export to
511 (None to select the one from the graph)
512 :param verbose: inserts prints
513 :param name: to overwrite onnx name
514 :param rename: rename the names to get shorter names
515 :param autopep_options: :epkg:`autopep8` options
516 :return: python code
518 The following example shows what a python code creating a graph
519 implementing the KMeans would look like.
521 .. runpython::
522 :showcode:
523 :process:
525 import numpy
526 from sklearn.cluster import KMeans
527 from mlprodict.onnx_conv import to_onnx
528 from mlprodict.onnx_tools.onnx_export import export2xop
530 X = numpy.arange(20).reshape(10, 2).astype(numpy.float32)
531 tr = KMeans(n_clusters=2)
532 tr.fit(X)
534 onx = to_onnx(tr, X, target_opset=14)
535 code = export2xop(onx)
537 print(code)
538 """
539 if isinstance(model_onnx, str):
540 model_onnx = onnx.load(model_onnx)
542 code = export_template(model_onnx, templates=get_xop_template(),
543 opset=opset, verbose=verbose, name=name,
544 rename=rename, use_onnx_tensor=True,
545 autopep_options=autopep_options)
546 return code