先把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)
具体过程参考:https://zhuanlan.zhihu.com/p/357844310
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中配置一下环境:
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)
其中指定的inpchannel输入通道要根据量化和训练测试的代码里面看:
此处是训练代码里面,opencv读入的是BGR,后面经过通道转换转成了RGB,再对网络进行输入,输入到网络,此处可以看出输入网络前的通道时RGB
又或者看测试代码:
此处的顺序是2 1 0,进行了通道转化,所以配置文件里面的指定输入通道时BGR
@@@@@@@@@@@@@@@@@@@@@@
此处有两种形式:
config量化 配置文件中的指定InpChannel为BGR→→→→→量化rknn时,进行通道转换2 1 0 ,BGR–>RGB→→→→→ (保持送进网络输入层的是RGB)
config量化 配置文件中的指定InpChannel为RGB→→→→→量化rknn时,进行通道转换时候保持不变0 1 2 ,还是RGB→→→→→ (保持送进网络输入层的是RGB)
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
其中需要记住的是:
opencv读入cv2.imread的是BGR,其他的image.open(PIL)\imshow等都是RGB