【RK3588-linux开发】3、训练yolov5并在RK3588上运行

首先要看看官方文档，对于PC端的开放（模型转换与模拟测试）阅读，toolkit下载:

Rockchip_User_Guide_RKNN_Toolkit2_CN-1.2.0.pdf

板载部署（C++/C）阅读其NPU的SDK开发文档，SDK下载:

Rockchip_RKNPU_User_Guide_RKNN_API_V1.2.0_CN.pdf

按照下面步骤训练自己的yolov5模型，并转化为RKNN模型，并再PC端模拟和板载部署运行。

1. 训练yolov5
下载RK推荐的yolov5进行训练，该仓库的改动如下:

（1）将common文件中激活层修改为ReLU，此外模型结构、训练、测试及其他操作均与原版本 Yolov5 一致。模型测试、导出时增添 rknn_mode 模式，导出对rknn友好型模型。
（2）导出模型可选择去掉尾端的permute层，从而兼容rknn_yolov5_demo的c++部署代码。
（3）导出模型使用 --rknn_mode 时候，默认将 大尺寸的 maxpool 等价替换成 多个 小尺寸的 maxpool，对计算结果无影响，但可以显著提升在 rknpu 上的推理速度。

训练和yolov5官方一致，但注意其基于的是较旧的版本！

2. torch训练的(pt)模型转onnx
调用models/export.py文件进行onnx的转化
参数修改:torch模型、 输入尺寸

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=r'models/320best_rkn.pt', help='weights path')  # from yolov5/models/
    parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')  # height, width
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
    parser.add_argument('--grid',default=False, action='store_true', help='export Detect() layer grid')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--rknn_mode', action='store_true', help='export rknn-friendly onnx model')
    parser.add_argument('--ignore_output_permute', action='store_true', help='export model without permute layer,which can be used for rknn_yolov5_demo c++ code')
    parser.add_argument('--add_image_preprocess_layer', action='store_true', help='add image preprocess layer, benefit for decreasing rknn_input_set time-cost')
    opt = parser.parse_args()

导出模型的部分，注意修改output_names，修改为output_names=['out378', 'out439', 'out500']，后面转为rknn时进行load_onnx需要指定名称。

        import onnx

        print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
        f = opt.weights.replace('.pt', '.onnx')  # filename
        torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
                          output_names=['out378', 'out439', 'out500'],
                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # size(1,3,640,640)
                                        'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)

导出rknn友好的模型，忽略尾端的permute层。

python3 models/export.py --rknn_mode  --ignore_output_permute

指定输入:

ONNX_MODEL = '../models/320best_rkn.onnx'
RKNN_MODEL = '320best_rkn_mmse.rknn'
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'
target_platform = "rk3588"
QUANTIZE_ON = True

3. onnx转rknn模型

PC端安装其工具，可以创建一个python3.6虚拟环境，直接在/home/jiang/Repositories/rknn/rknn-toolkit2-master/packages文件夹下

pip install rknn_toolkit2-1.2.0_f7bb160f-cp36-cp36m-linux_x86_64.whl -i https://pypi.douban.com/simple

这里使用的是rknn-toolkit2-master/examples/onnx/yolov5里面的test.py文件，包含后面PC端模拟测试。
指定输入:

ONNX_MODEL = '../models/320best_rkn.onnx'
RKNN_MODEL = '320best_rkn_mmse.rknn'
IMG_PATH = './bus.jpg'
DATASET = './dataset.txt'
target_platform = "rk3588"
QUANTIZE_ON = True
BOX_THESH = 0.5
NMS_THRESH = 0.3
IMG_SIZE = 320

修改你数据集聚类的先验框大小:

def yolov5_post_process(input_data):
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[13, 9], [23, 8], [20, 12], [32, 9], [26, 13],
               [28, 19], [35, 20], [61, 18], [42, 26]]

主函数中进行转化:

    # Create RKNN object 
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0,0,0]], 
                std_values=[[255,255,255]], 
                output_tensor_type='int8',
                target_platform = target_platform,
                quantized_algorithm = "mmse", #normal
                quantized_method = 'channel' ,#layer
                optimization_level = 1  #0  1  2  3
                )
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL, inputs = ['images'], 
                        outputs=['out378', 'out439', 'out500'] , 
                        input_size_list = [[3,320,320]])
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

对于 rknn.config参数设置(Rockchip_User_Guide_RKNN_Toolkit2_CN-1.2.0.pdf):

4. PC端模拟
rknn.init_runtime()设置输入默认为空就为模拟器运行。

    # Init runtime environment
    print('--> Init runtime environment')
    ret = rknn.init_runtime()
    # ret = rknn.init_runtime('rk3588',device_id='8f26b7fc2f55c8f2')
    if ret != 0:
        print('Init runtime environment failed!')
        exit(ret)
    print('done')

5. RK3588运行

• 环境准备
  安装linux-rga，下载

  $ mkdir build
  $ cd build
  $ cp ../cmake-linux.sh ./
  $ chmod +x ./cmake-linux.sh
  $ ./cmake-linux.sh
  

• 利用其SDK的rknpu2-master/examples/rknn_yolov5_demo进行运行。
  修改:

include/postprocess.h 中的OBJ_CLASS_NUM修改为自己类别数量；
src/postprocess.cc中修改label_list路径以及先验框大小。

postprocess.h

#define OBJ_NAME_MAX_SIZE 16
#define OBJ_NUMB_MAX_SIZE 64
#define OBJ_CLASS_NUM     80
#define NMS_THRESH        0.6
#define BOX_THRESH        0.5
#define PROP_BOX_SIZE     (5+OBJ_CLASS_NUM)

postprocess.cc

#define LABEL_NALE_TXT_PATH "./model/coco_80_labels_list.txt"

static char *labels[OBJ_CLASS_NUM];

const int anchor0[6] = {10, 13, 16, 30, 33, 23};
const int anchor1[6] = {30, 61, 62, 45, 59, 119};
const int anchor2[6] = {116, 90, 156, 198, 373, 326};

• 编译与运行

根据指定平台修改build-linux_RK3588.sh中的交叉编译器所在目录的路径TOOL_CHAIN，例如修改成:
export TOOL_CHAIN=/usr

./build-linux_linux_RK3588.sh
./build/build_linux_aarch64/rknn_yolov5_demo   model.rknn    test.jpg