40 lines
887 B
Bash
40 lines
887 B
Bash
|
#!/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."
|