diff --git a/scripts/build_trt_engines.sh b/scripts/build_trt_engines.sh index 781ec5b..0c7dae1 100755 --- a/scripts/build_trt_engines.sh +++ b/scripts/build_trt_engines.sh @@ -6,7 +6,11 @@ # Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp: # ArcFace : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112 (N = embed batch) # SCRFD : 1x3x640x640 (fixed; we letterbox to this) -# TransNetV2 : 1x100x27x48x3 (fixed; scene detector window), input tensor "input" +# TransNetV2 : 1x100x27x48x3 (fixed; scene detector window) +# +# Input tensor names are read from each ONNX model at runtime rather than +# hardcoded, since they differ between models (LVFace-B: "data", arcface_r18: +# "input", arcface_w600k_{r50,mbf}: "input.1"). # # These trtexec-built engines are *not* picked up by the ORT TRT EP cache — # ORT uses its own engine format. The point of this script is: @@ -27,24 +31,43 @@ SCENE_MODEL="${SCENE_MODEL:-$MODELS/transnetv2.onnx}" run() { echo "+ $*"; "$@"; } -echo "== ArcFace ==" +# The input tensor name is not the same across models — LVFace-B uses "data", +# arcface_r18 uses "input", and arcface_w600k_{r50,mbf} use "input.1". A +# hardcoded name makes trtexec fail with "Cannot find input tensor with name +# ...", so read it from the model instead. +input_name() { + python3 - "$1" <<'PY' +import sys +try: + import onnxruntime as ort +except ImportError: + sys.exit("onnxruntime is required to read the model's input name") +sess = ort.InferenceSession(sys.argv[1], providers=["CPUExecutionProvider"]) +print(sess.get_inputs()[0].name) +PY +} + +ARCFACE_IN="$(input_name "$ARCFACE_MODEL")" +SCRFD_IN="$(input_name "$SCRFD_MODEL")" + +echo "== ArcFace == (input tensor: $ARCFACE_IN)" run trtexec \ --onnx="$ARCFACE_MODEL" \ --fp16 \ - --minShapes=input.1:1x3x112x112 \ - --optShapes=input.1:${EMBED_BATCH}x3x112x112 \ - --maxShapes=input.1:${EMBED_BATCH}x3x112x112 \ + --minShapes="$ARCFACE_IN":1x3x112x112 \ + --optShapes="$ARCFACE_IN":${EMBED_BATCH}x3x112x112 \ + --maxShapes="$ARCFACE_IN":${EMBED_BATCH}x3x112x112 \ --saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \ --useCudaGraph echo -echo "== SCRFD ==" +echo "== SCRFD == (input tensor: $SCRFD_IN)" run trtexec \ --onnx="$SCRFD_MODEL" \ --fp16 \ - --minShapes=input.1:1x3x640x640 \ - --optShapes=input.1:1x3x640x640 \ - --maxShapes=input.1:1x3x640x640 \ + --minShapes="$SCRFD_IN":1x3x640x640 \ + --optShapes="$SCRFD_IN":1x3x640x640 \ + --maxShapes="$SCRFD_IN":1x3x640x640 \ --saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \ --useCudaGraph @@ -53,12 +76,14 @@ if [[ -f "$SCENE_MODEL" ]]; then echo "== TransNetV2 (scene detector) ==" # Fixed 1x100x27x48x3 window. The raw-TRT scene detector backend loads this # engine directly via --scene-detector-engine; the ORT-TRT EP builds its own. + SCENE_IN="$(input_name "$SCENE_MODEL")" + echo " (input tensor: $SCENE_IN)" run trtexec \ --onnx="$SCENE_MODEL" \ --fp16 \ - --minShapes=input:1x100x27x48x3 \ - --optShapes=input:1x100x27x48x3 \ - --maxShapes=input:1x100x27x48x3 \ + --minShapes="$SCENE_IN":1x100x27x48x3 \ + --optShapes="$SCENE_IN":1x100x27x48x3 \ + --maxShapes="$SCENE_IN":1x100x27x48x3 \ --saveEngine="$OUT/transnetv2.100x27x48.fp16.engine" \ --useCudaGraph else