fix(trt): read ONNX input tensor names instead of hardcoding input.1
build_trt_engines.sh hardcoded 'input.1' for the ArcFace and SCRFD shape
profiles, which only matches arcface_w600k_{r50,mbf}. Building engines for
any other embedder failed with:
Cannot find input tensor with name "input.1" in the network inputs!
Input names differ per model: LVFace-B_Glint360K uses 'data', arcface_r18
uses 'input', arcface_w600k_{r50,mbf} use 'input.1'. This matters now that
LVFace-B is the default embedder (src/config.hpp), so ARCFACE_MODEL=<LVFace>
is the expected path.
Read the name from each model via onnxruntime at build time.
This commit is contained in:
@@ -6,7 +6,11 @@
|
|||||||
# Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp:
|
# Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp:
|
||||||
# ArcFace : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112 (N = embed batch)
|
# ArcFace : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112 (N = embed batch)
|
||||||
# SCRFD : 1x3x640x640 (fixed; we letterbox to this)
|
# 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 —
|
# 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:
|
# 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 "+ $*"; "$@"; }
|
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 \
|
run trtexec \
|
||||||
--onnx="$ARCFACE_MODEL" \
|
--onnx="$ARCFACE_MODEL" \
|
||||||
--fp16 \
|
--fp16 \
|
||||||
--minShapes=input.1:1x3x112x112 \
|
--minShapes="$ARCFACE_IN":1x3x112x112 \
|
||||||
--optShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
--optShapes="$ARCFACE_IN":${EMBED_BATCH}x3x112x112 \
|
||||||
--maxShapes=input.1:${EMBED_BATCH}x3x112x112 \
|
--maxShapes="$ARCFACE_IN":${EMBED_BATCH}x3x112x112 \
|
||||||
--saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
|
--saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
|
||||||
--useCudaGraph
|
--useCudaGraph
|
||||||
|
|
||||||
echo
|
echo
|
||||||
echo "== SCRFD =="
|
echo "== SCRFD == (input tensor: $SCRFD_IN)"
|
||||||
run trtexec \
|
run trtexec \
|
||||||
--onnx="$SCRFD_MODEL" \
|
--onnx="$SCRFD_MODEL" \
|
||||||
--fp16 \
|
--fp16 \
|
||||||
--minShapes=input.1:1x3x640x640 \
|
--minShapes="$SCRFD_IN":1x3x640x640 \
|
||||||
--optShapes=input.1:1x3x640x640 \
|
--optShapes="$SCRFD_IN":1x3x640x640 \
|
||||||
--maxShapes=input.1:1x3x640x640 \
|
--maxShapes="$SCRFD_IN":1x3x640x640 \
|
||||||
--saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
|
--saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
|
||||||
--useCudaGraph
|
--useCudaGraph
|
||||||
|
|
||||||
@@ -53,12 +76,14 @@ if [[ -f "$SCENE_MODEL" ]]; then
|
|||||||
echo "== TransNetV2 (scene detector) =="
|
echo "== TransNetV2 (scene detector) =="
|
||||||
# Fixed 1x100x27x48x3 window. The raw-TRT scene detector backend loads this
|
# Fixed 1x100x27x48x3 window. The raw-TRT scene detector backend loads this
|
||||||
# engine directly via --scene-detector-engine; the ORT-TRT EP builds its own.
|
# 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 \
|
run trtexec \
|
||||||
--onnx="$SCENE_MODEL" \
|
--onnx="$SCENE_MODEL" \
|
||||||
--fp16 \
|
--fp16 \
|
||||||
--minShapes=input:1x100x27x48x3 \
|
--minShapes="$SCENE_IN":1x100x27x48x3 \
|
||||||
--optShapes=input:1x100x27x48x3 \
|
--optShapes="$SCENE_IN":1x100x27x48x3 \
|
||||||
--maxShapes=input:1x100x27x48x3 \
|
--maxShapes="$SCENE_IN":1x100x27x48x3 \
|
||||||
--saveEngine="$OUT/transnetv2.100x27x48.fp16.engine" \
|
--saveEngine="$OUT/transnetv2.100x27x48.fp16.engine" \
|
||||||
--useCudaGraph
|
--useCudaGraph
|
||||||
else
|
else
|
||||||
|
|||||||
Reference in New Issue
Block a user