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 DeepLabnum_classes = 2
input_shape = [512, 512]
backbone = "mobilenet"  # backbone：mobilenet、xception
downsample_factor = 16  # downsample factor, same with training
pretrained = Falsecuda = False    # use cuda
train = False  # model.train() mde
simplify = True  # simplify onnxmodel_path = 'model_data/22.pth'  # *.pth model pathdef 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) BCHWinput_layer_names = ["images"]output_layer_names = ["output"]# Update modelmodel.train() if train else model.eval()  # training mode = no Detect() layer grid construction# Export the modelprint(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)# Checksmodel_onnx = onnx.load(model_path)  # load onnx modelonnx.checker.check_model(model_onnx)  # check onnx model# Simplify onnxif simplify:import onnxsimprint(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 ImageONNX_MODEL = 'best.onnx'
RKNN_MODEL = 'best.rknn'
DATASET = './dataset.txt'if __name__ == '__main__':# Create RKNN objectrknn = RKNN(verbose=True)# pre-process configprint('--> config model')rknn.config(mean_values=[0, 0, 0], std_values=[255,255, 255], target_platform='rk3588')# Load modelprint('--> Loading model')ret = rknn.load_onnx(model=ONNX_MODEL)if ret != 0:print('Load model failed!')exit(ret)print('done')# Build modelprint('--> Building model')ret = rknn.build(do_quantization=True, dataset=DATASET)if ret != 0:print('Build model failed!')exit(ret)print('done')# Export RKNN modelprint('--> Export rknn model')ret = rknn.export_rknn(RKNN_MODEL)if ret != 0:print('Export rknn model failed!')exit(ret)print('done')# Init runtime environmentprint('--> 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+的测试的图片处理方式处理图片。

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