Skip to content

torch_to_nnef.inference_target.tract

Tools to manipulate tract programatically.

NOTE: interaction are done with *Nix tty system in mind, no support for Windows

Classes:

Name Description
TractBinaryDownloader

Tract Downloader.
TractCheckTolerance

Level of tolerated difference between output values of PyTorch and tract.
TractCli

tract calls from CLI.
TractNNEF

Tract NNEF inference target.

Functions:

Name Description
assert_io

Simple assertion without debug bundle.
assert_io_and_debug_bundle

Core check to ensure tract give same output as PyTorch within bounds.
debug_dumper_pytorch_to_onnx_to_nnef

Try to export the model with ONNX and convert the ONNX to NNEF via tract.

TractBinaryDownloader 

TractBinaryDownloader(version: SemanticVersion, auto_download: bool = True)

Tract Downloader.

NOTE: Current version assume you are using hardware officialy supported by tract with pre-built binaries.

Methods:

Name Description
dl_tract

Download tract requested version in cache directory.

Attributes:

Name Type Description
arch

Current OS architecture name needed to download tract cli asset.

arch property 

arch

Current OS architecture name needed to download tract cli asset.

dl_tract 

dl_tract()

Download tract requested version in cache directory.

TractCheckTolerance

Bases: str, Enum

Level of tolerated difference between output values of PyTorch and tract.

(those are defined in tract)

TractCli 

TractCli(tract_path: Path)

tract calls from CLI.

Why not use python package provided since few release of tract ?

  • we do not want to be coupled with a python lib as we declare version requested in API because this would lead to the need for an auto package download/import then rollback (since original environement may use another version)

Methods:

Name Description
assert_io_cmd_str

Assert a NNEF asset has outputs within tolerance bound with tract.
download

Download tract requested version in cache directory.
run_save_outputs_cmd_str

Run a NNEF asset and save tract's outputs to an NPZ bundle.

assert_io_cmd_str 

assert_io_cmd_str(nnef_path: Path, input_bundle_path: Path, output_bundle_path: Path, check_tolerance: TractCheckTolerance = TractCheckTolerance.EXACT)

Assert a NNEF asset has outputs within tolerance bound with tract.

download classmethod 

download(version: SemanticVersion) -> TractCli

Download tract requested version in cache directory.

run_save_outputs_cmd_str 

run_save_outputs_cmd_str(nnef_path: Path, input_bundle_path: Path, output_bundle_path: Path) -> T.List[str]

Run a NNEF asset and save tract's outputs to an NPZ bundle.

Returns the command line as a list of strings, suitable for subprocess.run. Used by the proptest comparator to do the Python-side, NaN-aware comparison that tract's strict --assert-output-bundle cannot express.

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.

assert_io 

assert_io(nnef_file_path: Path, tract_cli: TractCli, input_bundle_path: Path, output_bundle_path: Path, check_tolerance: TractCheckTolerance = TractCheckTolerance.EXACT)

Simple assertion without debug bundle.

With addition of gc of model once output is generated.

assert_io_and_debug_bundle 

assert_io_and_debug_bundle(model_info: UnfoldModelInfo, nnef_file_path: Path, tract_cli: TractCli, input_bundle_path: Path, output_bundle_path: Path, debug_bundle_path: Optional[Path] = None, check_tolerance: TractCheckTolerance = TractCheckTolerance.EXACT)

Core check to ensure tract give same output as PyTorch within bounds.

debug_dumper_pytorch_to_onnx_to_nnef 

debug_dumper_pytorch_to_onnx_to_nnef(model_info: UnfoldModelInfo, target_folder: Path, input_bundle_path: Path, tract_cli: TractCli, raise_export_error: bool = True) -> bool

Try to export the model with ONNX and convert the ONNX to NNEF via tract.

Used in debug bundle build (if it works, that's give a valuable reference, to debug T2N)