RK3588-NPU pytorch-image-models 模型编译测试
一.背景 二.操作步骤 1.下载依赖 2.创建容器 3.安装依赖 4.创建脚本 A.生成模型名列表 B.生成ONNX模型 C.生成RKNN模型 D.批量测试脚本
测试RK3588-NPU对https://github.com/huggingface/pytorch-image-models.git中模型的支持程度 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
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
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/
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_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
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
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
