一.背景

• 测试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_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