Files
scene-actor-extraction/scripts/build_trt_engines.sh
T
dtourolle d753062c6c Initial commit: scene-actor-extraction pipeline
Source (KPN++ pipeline nodes, ArcFace embedders, SCRFD/YuNet detectors,
gallery builder), build scripts, and eval artifacts.

- external/KPN as a git submodule (gitea.tourolle.paris/dtourolle/KPN)
- ONNX models tracked via Git LFS (models/*.onnx)
- generated outputs, TensorRT engines, reference repos, and media ignored
2026-06-12 15:29:01 +02:00

52 lines
1.8 KiB
Bash
Executable File
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env bash
# Pre-build TensorRT engines for ArcFace and SCRFD with the same shape profiles
# the runtime nodes use. First-run ORT engine builds take 3090 s per model and
# block the pipeline; this script does it offline so cold starts are instant.
#
# 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)
#
# 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:
# (a) sanity-check that the ONNX models build under TRT at all;
# (b) measure pure inference latency without ORT overhead.
# Run scene_analyze normally and ORT will populate ./trt_cache itself.
set -euo pipefail
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
MODELS="$ROOT/models"
OUT="$ROOT/trt_cache"
mkdir -p "$OUT"
EMBED_BATCH="${EMBED_BATCH:-4}"
ARCFACE_MODEL="${ARCFACE_MODEL:-$MODELS/arcface_w600k_r50.onnx}"
SCRFD_MODEL="${SCRFD_MODEL:-$MODELS/scrfd_500m_bnkps.onnx}"
run() { echo "+ $*"; "$@"; }
echo "== ArcFace =="
run trtexec \
--onnx="$ARCFACE_MODEL" \
--fp16 \
--minShapes=input.1:1x3x112x112 \
--optShapes=input.1:${EMBED_BATCH}x3x112x112 \
--maxShapes=input.1:${EMBED_BATCH}x3x112x112 \
--saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
--useCudaGraph
echo
echo "== SCRFD =="
run trtexec \
--onnx="$SCRFD_MODEL" \
--fp16 \
--minShapes=input.1:1x3x640x640 \
--optShapes=input.1:1x3x640x640 \
--maxShapes=input.1:1x3x640x640 \
--saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
--useCudaGraph
echo
echo "Engines saved under: $OUT"
echo "Look for 'mean: ... ms' in each section for per-call latency."