43 lines
1009 B
Bash
43 lines
1009 B
Bash
|
#!/usr/bin/env bash
|
|||
|
|
# ATC 模型转换脚本
|
|||
|
|
# 在高性能电脑（CANN 环境）上执行
|
|||
|
|
set -euo pipefail
|
|||
|
|
|
|||
|
|
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
|||
|
|
cd "$SCRIPT_DIR"
|
|||
|
|
|
|||
|
|
MODEL=${MODEL:-yolov8n}
|
|||
|
|
INPUT_SIZE=${INPUT_SIZE:-640}
|
|||
|
|
OUTPUT_TYPE=${OUTPUT_TYPE:-FP16}
|
|||
|
|
|
|||
|
|
echo "=========================================="
|
|||
|
|
echo " ATC ONNX -> OM 模型转换"
|
|||
|
|
echo "=========================================="
|
|||
|
|
echo "模型: $MODEL.onnx"
|
|||
|
|
echo "输入尺寸: ${INPUT_SIZE}x${INPUT_SIZE}"
|
|||
|
|
echo "输出精度: $OUTPUT_TYPE"
|
|||
|
|
echo ""
|
|||
|
|
|
|||
|
|
# 限制并行度，避免内存溢出
|
|||
|
|
export TBE_IMPL_MODE=AI_CORE
|
|||
|
|
export TE_PARALLEL_COMPILER=1
|
|||
|
|
|
|||
|
|
atc \
|
|||
|
|
--model=${MODEL}.onnx \
|
|||
|
|
--framework=5 \
|
|||
|
|
--output=${MODEL} \
|
|||
|
|
--soc_version=Ascend310B4 \
|
|||
|
|
--input_format=NCHW \
|
|||
|
|
--input_shape="images:1,3,${INPUT_SIZE},${INPUT_SIZE}" \
|
|||
|
|
--output_type=${OUTPUT_TYPE}
|
|||
|
|
|
|||
|
|
echo ""
|
|||
|
|
if [ -f ${MODEL}.om ]; then
|
|||
|
|
echo "转换成功!"
|
|||
|
|
echo "输出: ${MODEL}.om"
|
|||
|
|
ls -lh ${MODEL}.om
|
|||
|
|
else
|
|||
|
|
echo "转换失败，请检查日志"
|
|||
|
|
exit 1
|
|||
|
|
fi
|