torch_to_nnef.inference_target.base

Classes:

Name	Description
`InferenceTarget`	Base abstract class to implement a new inference engine 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: 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(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(nnef_graph: Graph, active_custom_extensions: List[str])

Get called just after PyTorch graph is parsed.

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(model: Module) -> T.Dict[str, str]

Optional custom fragments to pass.