TensorRT-Demo/yolo/build_int8_engines.sh

40 lines
887 B
Bash
Executable File

#!/bin/bash
# I use this script to build INT8 TensorRT engines for various yolov3 and
# yolov4 models.
set -e
models="yolov3-tiny-416 yolov3-608 yolov3-spp-608 yolov4-tiny-416 yolov4-608"
# make sure all needed files are present
for m in ${models}; do
if [[ ! -f ${m}.cfg ]]; then
echo "ERROR: cannot find the file ${m}.cfg"
exit 1
fi
if [[ ! -f ${m}.onnx ]]; then
echo "ERROR: cannot find the file ${m}.onnx"
exit 1
fi
done
# create symbolic links to cfg and onnx files
for m in ${models}; do
m_head=${m%-*}
m_tail=${m##*-}
ln -sf ${m}.cfg ${m_head}-int8-${m_tail}.cfg
ln -sf ${m}.onnx ${m_head}-int8-${m_tail}.onnx
done
# build TensorRT engines
for m in ${models}; do
m_head=${m%-*}
m_tail=${m##*-}
echo ; echo === ${m_head}-int8-${m_tail} === ; echo
python3 onnx_to_tensorrt.py --int8 -m ${m_head}-int8-${m_tail}
done
echo
echo "Done."