torch_to_nnef.inference_target

Targeted inference engine.

We mainly focus our effort to best support SONOS 'tract' inference engine.

Stricter Khronos NNEF specification mode also exist but is less extensively tested.

Modules:

Name	Description
base
khronos
tract	Tools to manipulate tract programatically.

Classes:

Name	Description
InferenceTarget	Base abstract class to implement a new inference engine target.
KhronosNNEF	Khronos Specification compliant NNEF asset build.
TractNNEF	Tract NNEF inference target.

InferenceTarget

InferenceTarget(version: Union[SemanticVersion, str], check_io: bool = False)

Base abstract class to implement a new inference engine target.

Init InferenceTarget.

Each inference engine is supposed to have at least a version and a way to check output given an input.

Methods:

Name	Description
post_export	Get called after NNEF model asset is generated.
post_trace	Get called just after PyTorch graph is parsed.
pre_trace	Get called just before PyTorch graph is traced.
specific_fragments	Optional custom fragments to pass.

Attributes:

Name	Type	Description
has_dynamic_axes	bool	Define if user request dynamic axes to be in the NNEF graph.

has_dynamic_axes property

has_dynamic_axes: bool

Define if user request dynamic axes to be in the NNEF graph.

Some inference engines may not support it hence False by default.

post_export

post_export(model_info: UnfoldModelInfo, nnef_graph: Graph, exported_filepath: Path, debug_bundle_path: Optional[Path] = None)

Get called after NNEF model asset is generated.

This is typically where check_io is effectively applied.

post_trace

post_trace(nnef_graph: Graph, active_custom_extensions: List[str])

Get called just after PyTorch graph is parsed.

pre_trace

pre_trace(model: Module, input_names: Optional[List[str]], output_names: Optional[List[str]])

Get called just before PyTorch graph is traced.

(after auto wrapper)

specific_fragments

specific_fragments(model: Module) -> T.Dict[str, str]

Optional custom fragments to pass.

KhronosNNEF

KhronosNNEF(version: Union[SemanticVersion, str], check_io: bool = True)

Bases: InferenceTarget

Khronos Specification compliant NNEF asset build.

in case of check_io=True we perform evaluation against nnef_tool nnef to pytorch converter. And access original and reloaded pytorch model provide same outputs

Methods:

Name	Description
post_export	Check io via the Torch interpreter of NNEF-Tools.

post_export

post_export(model_info: UnfoldModelInfo, nnef_graph: Graph, exported_filepath: Path, debug_bundle_path: Optional[Path] = None)

Check io via the Torch interpreter of NNEF-Tools.

TractNNEF

TractNNEF(version: Union[str, SemanticVersion], feature_flags: Optional[Set[TractFeatureFlag]] = None, check_io: bool = True, dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None, specific_tract_binary_path: Optional[Path] = None, check_io_tolerance: TractCheckTolerance = TractCheckTolerance.APPROXIMATE, specific_properties: Optional[Dict[str, str]] = None, dump_identity_properties: bool = True, force_attention_inner_in_f32: bool = False, force_linear_accumulation_in_f32: bool = False, force_norm_in_f32: bool = False, reify_sdpa_operator: Optional[bool] = None, upsample_with_debox: bool = False)

Bases: InferenceTarget

Tract NNEF inference target.

Init.

Parameters:

Name	Type	Description	Default
version	Union[str, SemanticVersion]	tract version targeted for export	required
feature_flags	Optional[Set[TractFeatureFlag]]	set of possibly added feature flags from tract (for example complex numbers)	None
check_io	bool	check between tract cli and Pytorch original model that given provided input, output is similar	True
dynamic_axes	Optional[Dict[str, Dict[int, str]]]	Optional specification of dynamic dimension By default the exported model will have the shapes of all input and output tensors set to exactly match those given in args. To specify axes of tensors as dynamic (i.e. known only at runtime) set dynamic_axes to a dict with schema: KEY (str): an input or output name. Each name must also be provided in input_names or output_names. VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a list, each element is an axis index.	None
specific_tract_binary_path	Optional[Path]	filepath of tract cli in case of custom non released version of tract (for testing purpose)	None
check_io_tolerance	TractCheckTolerance	TractCheckTolerance level of difference tolerance between original output values and those generated by tract (those are defined tract levels)	APPROXIMATE
specific_properties	Optional[Dict[str, str]]	custom tract_properties you wish to add inside NNEF asset (will be parsed by tract as metadata)	None
dump_identity_properties	bool	add tract_properties relative to user identity (host, username, OS...), helpfull for debug	True
force_attention_inner_in_f32	bool	control if attention should be forced as f32 inside (even if inputs are all f16), usefull for unstable networks like qwen2.5	False
force_linear_accumulation_in_f32	bool	usefull for f16 models to ensure that output of f16. f16 matmul become f32 accumulators.	False
force_norm_in_f32	bool	ensure that all normalization layers are in f32 whatever the original PyTorch modeling.	False
reify_sdpa_operator	Optional[bool]	(Optional) enable the conversion of scaled_dot_product_attention as a tract operator (intead of a NNEF fragment), default false until tract v0.22.0 included then true, except if specified. Experimental feature.	None
upsample_with_debox	bool	use debox upsample operator instead of deconvolution. This should be faster. (if tract version support it). Experimental feature.	False

Methods:

Name	Description
post_export	Perform check io and build debug bundle if fail.
post_trace	Add dynamic axes in the NNEF graph.
pre_trace	Check dynamic_axes are correctly formated.
specific_fragments	Optional custom fragments to pass.

post_export

post_export(model_info: UnfoldModelInfo, nnef_graph: Graph, exported_filepath: Path, debug_bundle_path: Optional[Path] = None)

Perform check io and build debug bundle if fail.

post_trace

post_trace(nnef_graph, active_custom_extensions)

Add dynamic axes in the NNEF graph.

pre_trace

pre_trace(model: Module, input_names: Optional[List[str]], output_names: Optional[List[str]])

Check dynamic_axes are correctly formated.

specific_fragments

specific_fragments(model: Module) -> T.Dict[str, str]

Optional custom fragments to pass.