tensorrt_exporter

sleap_nn.export.exporters.tensorrt_exporter

TensorRT export utilities.

Functions:

Name	Description
export_to_tensorrt	Export a PyTorch model to TensorRT format.

export_to_tensorrt(model, save_path, input_shape=(1, 1, 512, 512), input_dtype=torch.uint8, precision='fp16', min_shape=None, opt_shape=None, max_shape=None, workspace_size=2 << 30, method='onnx', verbose=True)

Export a PyTorch model to TensorRT format.

This function supports multiple compilation methods: - "onnx": Exports to ONNX first, then compiles with TensorRT (most reliable) - "jit": Uses torch.jit.trace + torch_tensorrt.compile (alternative)

Parameters:

Name	Type	Description	Default
model	Module	The PyTorch model to export (typically an ONNX wrapper).	required
save_path	str | Path	Path to save the TensorRT engine (.trt file).	required
input_shape	Tuple[int, int, int, int]	(B, C, H, W) optimal input tensor shape.	(1, 1, 512, 512)
input_dtype	dtype	Input tensor dtype (torch.uint8 or torch.float32).	uint8
precision	str	Model precision - "fp32" or "fp16".	'fp16'
min_shape	Optional[Tuple[int, int, int, int]]	Minimum input shape for dynamic shapes (default: batch=1, H/W halved).	None
opt_shape	Optional[Tuple[int, int, int, int]]	Optimal input shape (default: same as input_shape).	None
max_shape	Optional[Tuple[int, int, int, int]]	Maximum input shape (default: batch=16, H/W doubled).	None
workspace_size	int	TensorRT workspace size in bytes (default 2GB).	2 << 30
method	str	Compilation method - "onnx" or "jit".	'onnx'
verbose	bool	Print export info.	True

Returns:

Type	Description
Path	Path to the exported TensorRT engine.

Note

TensorRT models are NOT cross-platform. The exported model will only work on the same GPU architecture and TensorRT version used for export.

Source code in sleap_nn/export/exporters/tensorrt_exporter.py

def export_to_tensorrt(
    model: nn.Module,
    save_path: str | Path,
    input_shape: Tuple[int, int, int, int] = (1, 1, 512, 512),
    input_dtype: torch.dtype = torch.uint8,
    precision: str = "fp16",
    min_shape: Optional[Tuple[int, int, int, int]] = None,
    opt_shape: Optional[Tuple[int, int, int, int]] = None,
    max_shape: Optional[Tuple[int, int, int, int]] = None,
    workspace_size: int = 2 << 30,  # 2GB default
    method: str = "onnx",
    verbose: bool = True,
) -> Path:
    """Export a PyTorch model to TensorRT format.

    This function supports multiple compilation methods:
    - "onnx": Exports to ONNX first, then compiles with TensorRT (most reliable)
    - "jit": Uses torch.jit.trace + torch_tensorrt.compile (alternative)

    Args:
        model: The PyTorch model to export (typically an ONNX wrapper).
        save_path: Path to save the TensorRT engine (.trt file).
        input_shape: (B, C, H, W) optimal input tensor shape.
        input_dtype: Input tensor dtype (torch.uint8 or torch.float32).
        precision: Model precision - "fp32" or "fp16".
        min_shape: Minimum input shape for dynamic shapes (default: batch=1, H/W halved).
        opt_shape: Optimal input shape (default: same as input_shape).
        max_shape: Maximum input shape (default: batch=16, H/W doubled).
        workspace_size: TensorRT workspace size in bytes (default 2GB).
        method: Compilation method - "onnx" or "jit".
        verbose: Print export info.

    Returns:
        Path to the exported TensorRT engine.

    Note:
        TensorRT models are NOT cross-platform. The exported model will only
        work on the same GPU architecture and TensorRT version used for export.
    """
    import tensorrt as trt

    model.eval()
    device = next(model.parameters()).device

    save_path = Path(save_path)
    if not save_path.suffix:
        save_path = save_path.with_suffix(".trt")

    B, C, H, W = input_shape

    if min_shape is None:
        min_shape = (1, C, H // 2, W // 2)
    if opt_shape is None:
        opt_shape = input_shape
    if max_shape is None:
        max_shape = (min(16, B * 4), C, H * 2, W * 2)

    if verbose:
        print(f"Exporting model to TensorRT...")
        print(f"  Input shape: {input_shape}")
        print(f"  Min/Opt/Max: {min_shape} / {opt_shape} / {max_shape}")
        print(f"  Precision: {precision}")
        print(f"  Workspace: {workspace_size / 1e9:.1f} GB")
        print(f"  Method: {method}")

    if method == "onnx":
        return _export_tensorrt_onnx(
            model,
            save_path,
            input_shape,
            input_dtype,
            min_shape,
            opt_shape,
            max_shape,
            precision,
            workspace_size,
            verbose,
        )
    elif method == "jit":
        return _export_tensorrt_jit(
            model,
            save_path,
            input_shape,
            input_dtype,
            min_shape,
            opt_shape,
            max_shape,
            precision,
            workspace_size,
            verbose,
        )
    else:
        raise ValueError(f"Unknown method: {method}. Use 'onnx' or 'jit'.")