module utils.onnx_split
#
Short summary#
module onnxcustom.utils.onnx_split
Helpers to split an ONNX models.
Classes#
class |
truncated documentation |
---|---|
A segments of an onnx graph assuming it is the concatenation of all segments. |
|
The final goal is to split an onnx model into equivalent pieces. |
Functions#
function |
truncated documentation |
---|---|
Splits an ONNX model into n_parts consecutive subgraphs. Chained altogether, they are equivalent to the given model. … |
Static Methods#
staticmethod |
truncated documentation |
---|---|
Tells if a tensor is small. In that case, all edges to this constant are ignored when looking for cutting points. … |
Methods#
method |
truncated documentation |
---|---|
Builds one onnx subpart including segments from a to b (excluded). |
|
Splits the segments into two groups of the same size. |
|
Builds onnx subparts based on the segmentation defined by extremities. |
|
Splits the segments into n_parts segments |
Documentation#
Helpers to split an ONNX models.
- class onnxcustom.utils.onnx_split.OnnxSegment(parent, begin, end, size=0, involved=None, nodes=None)#
Bases:
object
A segments of an onnx graph assuming it is the concatenation of all segments.
- Parameters:
parent – an instance of OnnxSplitting
begin – names of the first extremity, None for the inputs of the main graph
end – names of the second extremity, None for the outputs of the main graph
size – total size of the segment
involved – list of result names involved in this segment
nodes – involved nodes, list of tuple (int, NodeProt)
- __init__(parent, begin, end, size=0, involved=None, nodes=None)#
- __repr__()#
Return repr(self).
- class onnxcustom.utils.onnx_split.OnnxSplitting(onnx_model, verbose=0, fLOG=None)#
Bases:
object
The final goal is to split an onnx model into equivalent pieces.
- Parameters:
onnx_model – onnx_model
verbose – displays information during the split
fLOG – logging function
- __init__(onnx_model, verbose=0, fLOG=None)#
- static _connex_components(vertices, adja)#
- _get_cutting_points(node_list)#
- _init()#
- static _key(idn, node)#
- _make_onnx(a, b, index=None)#
Builds one onnx subpart including segments from a to b (excluded).
- _make_segment(name1, name2)#
- _split_2(a, b)#
Splits the segments into two groups of the same size.
- Parameters:
a – first segment (included)
b – second segment (excluded)
- Returns:
split index
- static is_small(tensor)#
Tells if a tensor is small. In that case, all edges to this constant are ignored when looking for cutting points. The algorithm assumes it can be duplicated in multiple parts. It is usually single float constant or shapes.
- make_onnx(extremities)#
Builds onnx subparts based on the segmentation defined by extremities.
- Parameters:
extremities – example, [0, 3, 5], first onnx part contains segments 0:3=[0, 1, 2], second onnx part contains segments 3:5=[3, 4]
- Returns:
list of onnx subgraphs (:epkg:`ModelProto`)
- split_segment(n_parts=None, cut_points=None)#
Splits the segments into n_parts segments
- Parameters:
n_parts – number of parts to get
cut_points – uses this particular cut points
- Returns:
list of segments indices
- onnxcustom.utils.onnx_split.split_onnx(onnx_model, n_parts=None, cut_points=None, verbose=0, stats=False, fLOG=None)#
Splits an ONNX model into n_parts consecutive subgraphs. Chained altogether, they are equivalent to the given model.
- Parameters:
onnx_model – onnx model
n_parts – number of subgraphs
cut_points – use different cutting points that the ones the algorithm can found, it must be in the available set of cutting points
verbose – display information related to the split
stats – returns statistics as well, return of the function is a tuple
fLOG – logging function
- Returns:
list of onnx model