Coverage for mlprodict/cli/optimize.py : 100%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2@file
3@brief Command line about model optimisation.
4"""
5import os
6import onnx
9def onnx_stats(name, optim=False):
10 """
11 Computes statistics on an ONNX model.
13 :param name: filename
14 :param optim: computes statistics before an after optimisation was done
16 .. cmdref::
17 :title: Computes statistics on an ONNX graph
18 :cmd: -m mlprodict onnx_stats --help
19 :lid: l-cmd-onnx_stats
21 The command computes statistics on an ONNX model.
22 """
23 from ..onnx_tools.optim import onnx_statistics
24 if not os.path.exists(name):
25 raise FileNotFoundError( # pragma: no cover
26 "Unable to find file '{}'.".format(name))
27 with open(name, 'rb') as f:
28 model = onnx.load(f)
29 return onnx_statistics(model, optim=optim)
32def onnx_optim(name, outfile=None, recursive=True, options=None, verbose=0, fLOG=None):
33 """
34 Optimises an ONNX model.
36 :param name: filename
37 :param outfile: output filename
38 :param recursive: processes the main graph and the subgraphs
39 :param options: options, kind of optimize to do
40 :param verbose: display statistics before and after the optimisation
41 :param fLOG: logging function
43 .. cmdref::
44 :title: Optimises an ONNX graph
45 :cmd: -m mlprodict onnx_optim --help
46 :lid: l-cmd-onnx_optim
48 The command optimises an ONNX model.
49 """
50 from ..onnx_tools.optim import onnx_statistics, onnx_optimisations
51 if not os.path.exists(name):
52 raise FileNotFoundError( # pragma: no cover
53 "Unable to find file '{}'.".format(name))
54 if outfile == "":
55 outfile = None # pragma: no cover
56 if options == "":
57 options = None # pragma: no cover
58 if verbose >= 1 and fLOG is not None:
59 fLOG("[onnx_optim] read file '{}'.".format(name))
60 with open(name, 'rb') as f:
61 model = onnx.load(f)
62 if verbose >= 1 and fLOG is not None:
63 stats = onnx_statistics(model, optim=False)
64 for k, v in sorted(stats.items()):
65 fLOG(' before.{}={}'.format(k, v))
66 new_model = onnx_optimisations(model, recursive=recursive)
67 if verbose >= 1 and fLOG is not None:
68 stats = onnx_statistics(model, optim=False)
69 for k, v in sorted(stats.items()):
70 fLOG(' after.{}={}'.format(k, v))
71 if outfile is not None:
72 fLOG("[onnx_optim] write '{}'.".format(outfile))
73 with open(outfile, 'wb') as f:
74 onnx.save(new_model, f)
75 return new_model