6.1 deeplabv3+的pth模型转换为rknn模型

        和yolov5的pth模型转换为rknn模型类似，deeplabv3+的pth模型转为rknn模型的步骤是：

pth------>onnx-------->rknn

1.pth转为onnx

代码如下：

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# by [jackhanyuan](https://github.com/jackhanyuan)

import onnx
import torch
from pathlib import Path
from nets.deeplabv3_plus import DeepLab

num_classes = 2
input_shape = [512, 512]
backbone = "mobilenet"  # backbone：mobilenet、xception
downsample_factor = 16  # downsample factor, same with training
pretrained = False

cuda = False    # use cuda
train = False  # model.train() mde
simplify = True  # simplify onnxmodel_path = 'model_data/22.pth'  # *.pth model path


def convert_to_onnx(model_path):
    device = torch.device('cuda' if torch.cuda.is_available() and cuda else 'cpu')
    model = DeepLab(num_classes=num_classes, backbone=backbone, downsample_factor=downsample_factor,
                    pretrained=pretrained)
    model_path = Path(model_path)

    model.load_state_dict(torch.load(model_path, map_location=device))
    if cuda:
        model = model.cuda()
    print('{} loaded.'.format(model_path))

    im = torch.zeros(1, 3, *input_shape).to(device)  # image size(1, 3, 512, 512) BCHW
    input_layer_names = ["images"]
    output_layer_names = ["output"]

    # Update model
    model.train() if train else model.eval()  # training mode = no Detect() layer grid construction

    # Export the model
    print(f'Starting export with onnx {onnx.__version__}.')
    model_path = model_path.with_suffix('.onnx')
    torch.onnx.export(model,
                      im,
                      f=model_path,
                      verbose=False,
                      opset_version=12,
                      training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                      do_constant_folding=not train,
                      input_names=input_layer_names,
                      output_names=output_layer_names,
                      dynamic_axes=None)

    # Checks
    model_onnx = onnx.load(model_path)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model

    # Simplify onnx
    if simplify:
        import onnxsim
        print(f'Simplifying with onnx-simplifier {onnxsim.__version__}.')
        model_onnx, check = onnxsim.simplify(
            model_onnx,
            dynamic_input_shape=False,
            input_shapes=None)
        assert check, 'assert check failed'
        onnx.save(model_onnx, model_path)

    print('Onnx model save as {}'.format(model_path))


if __name__ == '__main__':
    convert_to_onnx('model_data/22.pth')#model_path

以上代码中需要修改以下内容：

1.num_classes = 2 #要检测的类的个数（要加个背景类，所以比实际类多1）;
2.input_shape = [512, 512]#要检测的图片的长宽（该模型的输入图片大小是512X512);
3.backbone = "mobilenet"  # backbone：mobilenet、xception(该模型用的卷积网络是mobilenet或者xception，训练时用哪个模型就在这里选哪个模型）；
4.downsample_factor = 16  # downsample factor, same with training（该模型的下采样用8或者16，训练时用哪个值在这里就用哪个值）；

5.convert_to_onnx('model_data/22.pth')#（要转换的训练好的pth模型路径）

注意：以上任何一个数值与训练时的数值不相同，则转换rknn模型时报错！！！！

2.onnx模型转换为rknn模型

模型转换的环境配置看我写的另一篇博文1.rk3588的yolov5运行：pt_onnx_rknn转换及rknn在rk3588系统python运行_rk3588支持的onnx版本-CSDN博客

模型转换代码如下：

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from PIL import Image


ONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'best.rknn'
DATASET = './dataset.txt'


if __name__ == '__main__':

    # 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], target_platform='rk3588')
    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=True, 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')

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

需要修改内容如下：

1.ONNX_MODEL = 'best.onnx'  #要转换的onnx模型；
2.RKNN_MODEL = 'best.rknn'  #转换后的rknn模型路径及名称；
3.DATASET = './dataset.txt'   #该文件里是要测的图片路径及名称；

注意：模型用到的图片大小必须是512X512X3的，如果是其他大小的图片，则需要处理转换成该大小的图片，图片的处理建议根据deeplabv3+的测试的图片处理方式处理图片。

在虚拟机上的代码运行结果如下：

        到这里，pth模型转换为rknn模型就完成了，后续将完成图片处理及rknn模型的检测。