pytorch模型转onnx-量化rknn（bisenet）

1.pytorch模型转化onnx

先把pytorch的.pth模型转成onnx,例如我这个是用Bisenet转的，执行export_onnx.py

import argparse
import os.path as osp
import sys
sys.path.insert(0, '.')

import torch

from lib.models import model_factory
from configs import set_cfg_from_file

torch.set_grad_enabled(False)


parse = argparse.ArgumentParser()
parse.add_argument('--config', dest='config', type=str,
        default='G:/6666Ground_segmentation0813/configs/bisenetv2_city.py',)
parse.add_argument('--weight-path', dest='weight_pth', type=str,
        default='G:/6666Ground_segmentation0813/v4_model_final.pth') #最后的pytorch模型
parse.add_argument('--outpath', dest='out_pth', type=str,
        default='G:/6666Ground_segmentation0813/model0124.onnx')#转成onnx的路径
args = parse.parse_args()


cfg = set_cfg_from_file(args.config)
if cfg.use_sync_bn: cfg.use_sync_bn = False

net = model_factory[cfg.model_type](cfg.n_cats, aux_mode='pred')
net.load_state_dict(torch.load(args.weight_pth), strict=False)
net.eval()


#  dummy_input = torch.randn(1, 3, *cfg.crop_size)
# dummy_input = torch.randn(1, 3, 1024, 2048)
dummy_input = torch.randn(1, 3, 480, 640) #图像的输入尺寸
input_names = ['input_image']
output_names = ['preds',]

torch.onnx.export(net, dummy_input, args.out_pth,
    input_names=input_names, output_names=output_names,
    verbose=False, opset_version=11)

转好之后长这样~

2.onnx转rknn量化模型

2.1 配置rknn环境（win64）

2.1.1 安装python3.6的环境

具体过程参考：https://zhuanlan.zhihu.com/p/357844310

2.1.2 安装依赖包

conda info --envs 查看环境
创建新的rknn避免和其他的环境产生冲突
1.conda create -n rknn python=3.6
2.activate rknn
3.python -m pip install --upgrade pip
4.pip install tensorflow==1.14.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
5.pip install torch==1.6.0+cpu torchvision==0.7.0+cpu -f https://download.pytorch.org/whl/torch_stable.html --user
6.pip install opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple
7.pip install G:\6666Ground_segmentation0813\rknn_toolkit-1.6.0-cp36-cp36m-win_amd64.whl
这里是rknn套件安装包路径

8.pip install mxnet==1.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
9.pip install tensorflow==1.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
下面这两个8和9貌似可以不安装，无伤大雅~ 中间如果版本冲突的话直接uninstall卸载再重装需要的版本，我是最后整了几次卸载掉matplotlib就好了

一顿操作之后：检查下python3.6版本问题是否符合要求

rknn套件安装包长这样：

在pycharm中配置一下环境：

2.2 量化的话需要准备量化校正数据集

这个txt里面放图片路径：

生成路径文件的python:

import os
# 图片文件夹路径
pic_paths= "G:/6666Ground_segmentation0813/rknn_dataset"
f=open('G:/6666Ground_segmentation0813/rknntxt.txt', 'w')
filenames=os.listdir(pic_paths)
filenames.sort()
for filename in filenames:
     # 图片绝对路径
    out_path="G:/6666Ground_segmentation0813/rknn_dataset/" + filename
    print(out_path)
    f.write(out_path+'\n')
f.close()

其中需要计算训练代码时候的里面的均值和归一化的值:

得到如下量化参数配置均值和归一化的值：

再用用转好的onnx模型进行量化，转成rknn模型:

from __future__ import absolute_import, print_function, division
import os
from rknn.api import RKNN

# onnx_model = './resource/onnx/model_256x256_max_mscf_0.924553.onnx'G:/6666Ground_segmentation0813
onnx_model = 'G:/6666Ground_segmentation0813/onnx/model0124.onnx' #onnx路径
save_rknn_dir = 'G:/6666Ground_segmentation0813/rknn'#rknn保存路径

if __name__ == '__main__':

    # Create RKNN object
    rknn = RKNN()

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[83.0535, 94.095, 82.1865]], std_values=[[53.856, 54.774, 53.9325]], reorder_channel='2 1 0', target_platform=['rk1808'], batch_size=1,quantized_dtype='asymmetric_quantized-u8')  # 需要输入为RGB#####需要转化一下均值和归一化的值
    # rknn.config(mean_values=[[0.0, 0.0, 0.0]], std_values=[[255, 255, 255]], reorder_channel='2 1 0', target_platform=['rv1126'], batch_size=1)  # 需要输入为RGB
    print('done')

    model_name = onnx_model[onnx_model.rfind('/') + 1:]
    # Load ONNX model
    print('--> Loading model %s' % model_name)
    ret = rknn.load_onnx(model=onnx_model)
    if ret != 0:
        print('Load %s failed!' % model_name)
        exit(ret)
    print('done')
    # Build model
    print('--> Building model')
    # ret = rknn.build(do_quantization=False, dataset='./quantization_dataset.txt', pre_compile=True) ###路哥的版本pre_compil=True离线预编译
    ret = rknn.build(do_quantization=True, dataset='G:/6666Ground_segmentation0813/rknntxt.txt', pre_compile=False)
    #do_quantization是否对模型进行量化，datase量化校正数据集，pre_compil模型预编译开关，预编译 RKNN 模型可以减少模型初始化时间，但是无法通过模拟器进行推理或性能评估
    if ret != 0:
        print('Build net failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    # save_name = model_name.replace(os.path.splitext(model_name)[-1], "_no_quant.rknn")
    save_name = model_name.replace(os.path.splitext(model_name)[-1], "_quant.rknn")
    ret = rknn.export_rknn(os.path.join(save_rknn_dir, save_name))
    if ret != 0:
        print('Export rknn failed!')
        exit(ret)
    print('done')

    rknn.release()

200张图片数据进行量化，大概11分钟左右（640*480），然后量化成功！

最后Pytorch的20M模型——转成onnx(12.7M)——量化RKNN(3.3M)

配置文件model_config.txt

其中指定的inpchannel输入通道要根据量化和训练测试的代码里面看：
此处是训练代码里面，opencv读入的是BGR,后面经过通道转换转成了RGB,再对网络进行输入，输入到网络，此处可以看出输入网络前的通道时RGB

又或者看测试代码：


此处的顺序是2 1 0,进行了通道转化，所以配置文件里面的指定输入通道时BGR
@@@@@@@@@@@@@@@@@@@@@@
此处有两种形式：

1. config量化 配置文件中的指定InpChannel为BGR→→→→→量化rknn时，进行通道转换2 1 0 ，BGR–>RGB→→→→→ (保持送进网络输入层的是RGB)

2. config量化 配置文件中的指定InpChannel为RGB→→→→→量化rknn时，进行通道转换时候保持不变0 1 2 ，还是RGB→→→→→ (保持送进网络输入层的是RGB)

※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
其中需要记住的是：
opencv读入cv2.imread的是BGR,其他的image.open(PIL)\imshow等都是RGB