onnx_extended.tools.onnx_io#

enumerate_model_tensors#

onnx_extended.tools.enumerate_model_tensors(model: ModelProto) → Iterable[Tuple[TensorProto, bool]][source]#

Enumerates all tensors in a model.

Parameters:: model – model to process
Returns:: iterator on a couple (TensorProto, bool), the boolean indicates if the data is external

load_external#

onnx_extended.tools.load_external(model: ModelProto, base_dir: str, names: Set[str] | None = None)[source]#

Loads external data into memory.

Parameters:

model – the model loaded with load_model()
base_dir – directory when the data can be found
names – subsets of names to load or None for all

load_model#

Loads a model or returns the only argument if the type is already a ModelProto.

Parameters:

model – proto file
external – loads the external data as well
base_dir – needed if external is True and the model has external weights

Returns:

ModelProto

onnx2string#

onnx_extended.tools.onnx_io.onnx2string(proto: ModelProto, as_code: bool = False) → str[source]#

Takes a model and returns a string pluggagle in a python script. It uses module base64.

Parameters:

proto – model proto
as_code – if true, returns the model as a piece of code

Returns:

string

save_model#

onnx_extended.tools.save_model(proto: ModelProto, filename: str, external: bool = False, convert_attribute: bool = True, size_threshold: int = 1024, all_tensors_to_one_file: bool = True)[source]#

Saves a model into an onnx file.

Parameters:

proto – ModelProto
filename – where to save it
external – saves weights as external data
convert_attribute – converts attributes as well
size_threshold – every weight above that threshold is saved as external
all_tensors_to_one_file – saves all tensors in one unique file

string2onnx#

onnx_extended.tools.onnx_io.string2onnx(text: str) → ModelProto[source]#

Restores a model saved with onnx2string.

Parameters:: text – string produced by onnx2string
Returns:: ModelProto