TensorRT-Demo/yolo/build_dla_engines.sh

44 lines
1.1 KiB
Bash
Raw Normal View History

2023-03-06 20:44:29 +08:00
#!/bin/bash
# I use this script to build DLA0 and DLA1 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}-dla0-${m_tail}.cfg
ln -sf ${m}.onnx ${m_head}-dla0-${m_tail}.onnx
ln -sf ${m}.cfg ${m_head}-dla1-${m_tail}.cfg
ln -sf ${m}.onnx ${m_head}-dla1-${m_tail}.onnx
done
# build TensorRT engines
for m in ${models}; do
m_head=${m%-*}
m_tail=${m##*-}
echo ; echo === ${m_head}-dla0-${m_tail} === ; echo
python3 onnx_to_tensorrt.py --int8 --dla_core 0 -m ${m_head}-dla0-${m_tail}
echo ; echo === ${m_head}-dla1-${m_tail} === ; echo
python3 onnx_to_tensorrt.py --int8 --dla_core 1 -m ${m_head}-dla1-${m_tail}
done
echo
echo "Done."