cat > list_models.py <<-'EOF'
from timm.models import create_model, is_model, list_models
model_names = list_models()
with open("models.txt","w") as f:
for n in model_names:
f.write(f"{n}\n")
EOF
python list_models.py
cat> gen_rknn.py<<-'EOF'
import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
ONNX_MODEL = 'model.onnx'
RKNN_MODEL = 'model.rknn'
if not os.path.exists(ONNX_MODEL):
exit(-1)
if os.path.exists(RKNN_MODEL):
os.remove(RKNN_MODEL)
if not os.path.exists("./dataset.txt"):
img=np.ones((224,224,3),dtype=np.int8)
cv2.imwrite("img.jpg",img)
with open("./dataset.txt","w") as f:
f.write("img.jpg")
rknn = RKNN(verbose=False)
rknn.config(mean_values=[123.675, 116.28, 103.53],
std_values=[58.82, 58.82, 58.82],
target_platform='rk3588')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
exit(ret)
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
exit(ret)
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
exit(ret)
rknn.release()
if os.path.exists(RKNN_MODEL):
print("BUILD SUCCESS")
EOF
D.批量测试脚本
bash复制代码
cat> run.sh<<-'EOF'
#!/bin/bash
mkdir -p ./log
# 获取 models.txt 中的模型总数
total=$(wc -l < models.txt)
count=0
for name in $(cat models.txt); do
if [ ! -f ./log/$name ]; then
# 打印进度条
progress=$((count * 100 / total))
echo -ne "进度: ["
for ((i = 0; i < progress; i+=2)); do echo -n '='; done
for ((i = progress; i < 100; i+=2)); do echo -n ' '; done
echo -ne "] $progress%\r"
echo $name
# 获取开始时间
model_start_time=$(date +%s)
# 运行python脚本
python gen_onnx.py $name > ./log/$name 2>&1
python gen_rknn.py >> ./log/$name 2>&1
# 记录结束时间并计算耗时
model_end_time=$(date +%s)
model_duration=$((model_end_time - model_start_time))
output=$(grep "BUILD SUCCESS" ./log/$name)
echo "$name $output (耗时: ${model_duration}s)"
# 增加已完成的模型计数
count=$((count + 1))
fi
done
# 打印新行以结束进度条
echo ""
EOF
bash run.sh