model_optimization/convert_tensorrt.py

"""TensorRT conversion utility wrapper."""

from __future__ import annotations

import shutil
import subprocess
from pathlib import Path


def convert(onnx: Path, engine: Path) -> None:
    """Convert ONNX to TensorRT engine when trtexec exists; fallback to copy."""
    engine.parent.mkdir(parents=True, exist_ok=True)
    trtexec = shutil.which("trtexec")
    if trtexec:
        subprocess.check_call([trtexec, f"--onnx={onnx}", f"--saveEngine={engine}", "--fp16"])
        return
    shutil.copyfile(onnx, engine)


if __name__ == "__main__":
    convert(Path("artifacts/models/iresnet100_arcface.onnx"), Path("artifacts/models/iresnet100_arcface.plan"))