Shortcuts#

IO#

enumerate_model_tensors#

onnx_extended.tools.enumerate_model_tensors(model: ModelProto) → Iterator[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#

onnx_extended.tools.load_model(model: str | ModelProto | GraphProto | FunctionProto, external: bool = True, base_dir: str | None = None) → ModelProto | GraphProto | FunctionProto[source]#

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

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