trtutils.trtexec package

Module contents

Submodule for utilities related to the trtexec tool.

Functions

build_engine()

Build an engine from an ONNX file using trtexec.

find_trtexec()

Find an instance of the trtexec binary on the system.

run_trtexec()

Run trtexec command.

trtutils.trtexec.build_engine(weights: Path, output: Path, use_dla_core: int | None = None, shapes: Sequence[tuple[str, tuple[int, ...]]] | None = None, workspace: int | None = None, add_args: Sequence[str] | None = None, *, fp16: bool | None = None, int8: bool | None = None, fp8: bool | None = None, best: bool | None = None, allow_gpu_fallback: bool | None = None) → bool

Build an engine from a weight file using trtexec.

Parameters:

• weights (Path) – The path to the weight file to build the engine from. Examples are: .onnx, .prototxt If a .onnx file is provided, the engine will be built from the ONNX model. If a .prototxt file is provided, the engine will be built with random weights based on the model architecture.

• output (Path) – The path to save the built engine to.

• use_dla_core (int, optional) – The DLA core to use for building the engine, by default None. The DLA core should be either 0 or 1 if specified.

• shapes (tuple[tuple[int, ...], ...], optional) – The input shapes to use for the engine, by default None. If provided, the engine will be built with these input shapes. The name of the input must also be defined. An example could be: ((“images”, (1, 3, 640, 640)), …)

• workspace (int, optional) – The workspace size to use for the engine, by default None. Expressed in MiB.

• fp16 (bool, optional) – Whether to use FP16 precision for the engine, by default None.

• int8 (bool, optional) – Whether to use INT8 precision for the engine, by default None.

• fp8 (bool, optional) – Whether to use FP8 precision for the engine, by default None.

• best (bool, optional) – Whether to use the best precision available for the engine, by default None.

• allow_gpu_fallback (bool, optional) – Whether to allow GPU fallback when a layer is not supported on DLA. By default, this is None.

• add_args (Sequence[str], optional) – Additional arguments to pass to trtexec, by default None.

Returns:

Whether the engine was built successfully.

Return type:

bool

Raises:

• FileNotFoundError – If the weight file is not found.

• IsADirectoryError – If the weight file is a directory.

• ValueError – If the weight file does not have a valid extension.

• ValueError – If the DLA core is not 0 or 1.

• RuntimeError – If the command generation failed.

• TypeError – If the input shapes are not integers.

trtutils.trtexec.find_trtexec() → Path

Find an instance of the trtexec binary on the system.

Requires the locate command to be installed on the system. As such, only works on Unix-like systems.

Returns:

The path to the trtexec binary

Return type:

Path

Raises:

FileNotFoundError – If the trtexec binary is not found on the system

trtutils.trtexec.run_trtexec(command: str, trtexec_path: Path | str | None = None) → tuple[bool, str, str]

Run a command using trtexec.

The goal of this function is make it easier to use trtexec within Python scripts. By returning the stdout/stderr streams via strings back to the Python program it can simplify logic or scripts which utilize trtexec.

Parameters:

• command (str) – The command to run using trtexec

• trtexec_path (Path | str | None, optional) – The path to the trtexec binary to use. If None, find_trtexec will be used.

Returns:

A tuple containing the following elements: (success, stdout, stderr)

Return type:

tuple[bool, str, str]