RK3588-NPU pytorch-image-models 模型编译测试

一.背景

• 测试RK3588-NPU对https://github.com/huggingface/pytorch-image-models.git中模型的支持程度

二.操作步骤

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_exportONNX_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 RKNNONNX_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/bashmkdir -p ./log
# 获取 models.txt 中的模型总数
total=$(wc -l < models.txt)
count=0
for name in $(cat models.txt); doif [ ! -f ./log/$name ]; then# 打印进度条progress=$((count * 100 / total))echo -ne "进度: ["for ((i = 0; i < progress; i+=2)); do echo -n '='; donefor ((i = progress; i < 100; i+=2)); do echo -n ' '; doneecho -ne "] $progress%\r"	echo $name        # 获取开始时间model_start_time=$(date +%s)        # 运行python脚本python gen_onnx.py $name > ./log/$name 2>&1python 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