module utils.onnx_split#

Inheritance diagram of onnxcustom.utils.onnx_split

Short summary#

module onnxcustom.utils.onnx_split

Helpers to split an ONNX models.

source on GitHub

Classes#

class

truncated documentation

OnnxSegment

A segments of an onnx graph assuming it is the concatenation of all segments.

OnnxSplitting

The final goal is to split an onnx model into equivalent pieces.

Functions#

function

truncated documentation

split_onnx

Splits an ONNX model into n_parts consecutive subgraphs. Chained altogether, they are equivalent to the given model. …

Static Methods#

staticmethod

truncated documentation

_connex_components

_key

is_small

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

__init__

__init__

__repr__

_get_cutting_points

_init

_make_onnx

Builds one onnx subpart including segments from a to b (excluded).

_make_segment

_split_2

Splits the segments into two groups of the same size.

make_onnx

Builds onnx subparts based on the segmentation defined by extremities.

split_segment

Splits the segments into n_parts segments

Documentation#

Helpers to split an ONNX models.

source on GitHub

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)

source on GitHub

__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

source on GitHub

__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).

source on GitHub

_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

source on GitHub

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.

source on GitHub

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`)

source on GitHub

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

source on GitHub

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

source on GitHub