Model (aka Typed Model)

class tract.model.Model(ptr)[source]

Main model object, central focus point of the model transformation pipeline.

The Model is the central point of tract model loading and “model cooking”. ONNX and NNEF serialized models are converted to Model (more or less directly) before we can do anything of value with them. Model can be dumped to NNEF (or tract-opl which is NNEF plus tract proprietary extensions).

A Model can be optimize()-d, substituting the “high level” operators in tract-core operator set by the best implementation available for the current system. From there it can be transformed into a Runnable object that we will use to run.

Model cooking

Some model transformations can be performed on the Model class:

  • declutter (getting rid of training artefacts)

  • “pulsification” (transforming a batch-oriented model into a streaming model)

  • symbol substitution (make N or Batch a fixed number, unlocking potential optimisation later on)

  • static cost evaluation and dynamic profiling

In some situations, these operations are done “on-the-fly” when an ONNX or NNEF model is loaded, at start-up time. In other situations, when start-up time becomes an issue, it may be beneficial to “pre-cook” the model: apply the transformations once, serialize the model as NNEF (with tract-opl extension if needed). At start-up this model can be significantly less expensive to “cook” for inference.

Model and TypedModel

This class is actually a wrapper around the “TypedModel” in Rust codebase. The “Typed” bit means that all shapes and element types in all input, output and temporary values must be known. There is support in tract for symbols in dimensions, with some limited computation capabilities on symbolic expressions. For instance, it is relatively frequent to work with a Model where all tensor shapes start with N or Batch.

input_count() int[source]

Return the number of inputs of the model

output_count() int[source]

Return the number of outputs of the model

input_name(input_id: int) str[source]

Return the name of the input_id-th input

input_fact(input_id: int) Fact[source]

Return the fact of the input_id-th input

set_output_names(names: List[str])[source]

Change the output nodes of the model

output_name(output_id: int) str[source]

Return the name of the output_id-th output

output_fact(output_id: int) Fact[source]

Return the fact of the output_id-th output

transform(transform: str | TransformSpec) None[source]

Apply a transform to the model.

transform can be:

  • a plain string name (e.g. "f32_to_f16")

  • a JSON string with a "name" key and parameters

  • a TransformSpec subclass such as ConcretizeSymbols or Pulse

into_runnable() Runnable[source]

Transform the model into a ready to be used Runnable model

property_keys() List[str][source]

Query the list of properties names of the model.

property(name: str) Tensor[source]

Query a property by name