master
c69d9b1457
- yolov8n.onnx (13MB) exported from ultralytics YOLOv8n - scripts/atc_convert.sh for ONNX->OM conversion on Ascend 310B4 - Use FP16 precision for faster inference on edge device
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
|