RK3588-NPU pytorch-image-models 模型编译测试
一.背景
二.操作步骤
1.下载依赖
mkdir rk3588
cd rk3588
wget -O v2.3.0.tar.gz https://github.com/airockchip/rknn-toolkit2/archive/refs/tags/v2.3.0.tar.gz
tar -xf v2.3.0.tar.gz
git clone https://github.com/huggingface/pytorch-image-models.git
2.创建容器
docker stop rk3588_npu_test
docker rm rk3588_npu_test
docker run --gpus all --shm-size=128g -id -e NVIDIA_VISIBLE_DEVICES=all \
--privileged --net=host -v $PWD:/home -w /home \
--name=rk3588_npu_test nvcr.io/nvidia/pytorch:23.07-py3 /bin/bash
docker start rk3588_npu_test
docker exec -ti rk3588_npu_test bash
3.安装依赖
cd /home
pip install -r rknn-toolkit2-2.3.0/rknn-toolkit2/packages/x86_64/requirements_cp310-2.3.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
pip install rknn-toolkit2-2.3.0/rknn-toolkit2/packages/x86_64/rknn_toolkit2-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
cd /home/pytorch-image-models
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
pip uninstall opencv-contrib-python opencv-python -y
pip install opencv-python==4.7.0.72 -i https://pypi.tuna.tsinghua.edu.cn/simple/
pip install opencv-contrib-python==4.7.0.72 -i https://pypi.tuna.tsinghua.edu.cn/simple/
4.创建脚本
A.生成模型名列表
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
B.生成ONNX模型
cat> gen_onnx.py<<-'EOF'
import timm
import sys
import torch
import os
from timm.utils.model import reparameterize_model
from timm.utils.onnx import onnx_export
ONNX_MODEL = 'model.onnx'
if os.path.exists(ONNX_MODEL):
os.remove(ONNX_MODEL)
model = timm.create_model(
sys.argv[1],
num_classes=2,
in_chans=3,
pretrained=False,
exportable=True,
)
model.eval()
input_size=model.default_cfg.get('input_size')
input_tensor = torch.ones((1,) + input_size)
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, input_tensor, ONNX_MODEL,
verbose=False, input_names=input_names,
output_names=output_names,opset_version=17,export_params=True)
EOF
C.生成RKNN模型
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.批量测试脚本
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