瑞芯微(EASY EAI)RV1126B AI超分辨率算法教程

烟火之旅 2026-09-22 1234

默认

1.PIPNet简介

XLSR作为2021年移动AI实时单图像超分辨率挑战赛而设计的一种新颖的图像超分辨率算法。它通过将高效移动网络模块设计理念与针对硬件特性的量化鲁棒性训练策略相结合，最终在移动设备上实现了一个兼具卓越重建质量、极低参数量和实时性能的轻量级超分网络，并赢得了该挑战赛的冠军。XLSR质量评估指标：PSNR值为29.58，SSIM值为0.86。

1.1 模型导出优化

为了使得XLSR算法在超分后细节更突出，在原模型的最后一层加了拉普拉斯算子层SharpeningLayer()，并且拉普拉斯算子核是可调的。同时，为了避免在RKNN模型后处理时对输出数据进行转置（NCHW转NHWC），在模型前向计算的最后一层加了转置层，达到加速目的。

在PC端执行export_onnx.py将pt模型转成onnx。如下代码段所示：

from model.models import * import torch h, w = 360, 640 scaling_factor = 2 # 2或4 weight = r"weight/xlsr_2x_checkpoint_float32.pth.tar" output_file = r"weight/xlsr_{}x_{}x{}.onnx".format(scaling_factor, h, w) model = XLSRRelease(scaling_factor=scaling_factor) checkpoint = torch.load(weight, map_location='cpu', weights_only=False) model.load_state_dict(checkpoint["state_dict"], strict=False) model.eval() input_size=(3, h, w) # c, h, w example_input = torch.randn((1,) + input_size, requires_grad=False) torch_out = torch.onnx.export( model, example_input, output_file, export_params=True, input_names=["image"], output_names=["upscaled_image"], opset_version=11 ) print("export finish")

2. 资料下载

进行后续操作前必须先下载完整资料包！！！

资料名称

链接

训练代码github

https://github.com/EASY-EAI/yolov5

超分辨率算法教程完整源码包

https://pan.baidu.com/s/1MWnMdrKDZl_4leHoKwAIyw?pwd=1234

提取码: 1234

3. rknn-toolkit模型转换

3.1 rknn-toolkit模型转换环境搭建

onnx模型需要转换为rknn模型才能在EASY-EAI-Nano-TB运行，所以需要先搭建rknn-toolkit模型转换工具的环境。当然tensorflow、tensroflow lite、caffe、darknet等也是通过类似的方法进行模型转换，只是本教程onnx为例。

3.1.1 概述

模型转换环境搭建流程如下所示：

3.1.2 下载模型转换工具

为了保证模型转换工具顺利运行，请在官网产品简介下载网盘里“06.AI算法开发/01.rknn-toolkit2模型转换工具/rknn-toolkit2-v2.3.0/rknn-toolkit2/docker/docker_image/rknn-toolkit2-v2.3.2-cp38-docker.tar.gz”。

网盘下载链接:https://pan.baidu.com/s/1cyvdodUCtSvWZsyIfqzQ-w?pwd=1234提取码: 1234

3.1.3 把工具移到ubuntu20.04

把下载完成的docker镜像移到我司的虚拟机ubuntu20.04的rknn-toolkit2目录,如下图所示：

3.1.4 运行模型转换工具环境

在该目录打开终端

执行以下指令加载模型转换工具docker镜像:

docker load --input rknn-toolkit2-v2.3.2-cp38-docker.tar.gz

执行以下指令进入镜像bash环境:

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb rknn-toolkit2:2.3.2-cp38 /bin/bash

现象如下图所示:

输入“python”加载python相关库，尝试加载rknn库，如下图环境测试成功:

至此，模型转换工具环境搭建完成。

4. 模型转换为RKNN

EASY-EAI-Nano-TB支持.rknn后缀的模型的评估及运行，对于常见的tensorflow、tensroflow lite、caffe、darknet、onnx和Pytorch模型都可以通过我们提供的 toolkit 工具将其转换至 rknn 模型，而对于其他框架训练出来的模型，也可以先将其转至 onnx 模型再转换为 rknn 模型。模型转换操作流程入下图所示：

4.1 模型转换Demo下载

打开目录《3.AI模型转换》把 xlsr-python-2x.tar 解压到虚拟机，如下图所示：

4.2 进入模型转换工具docker环境

执行以下指令把工作区域映射进docker镜像，其中/home/developer/rknn-toolkit2/model_convert为工作区域，/test为映射到docker镜像，/dev/bus/usb:/dev/bus/usb为映射usb到docker镜像：

docker run -t -i --privileged -v /dev/bus/usb:/dev/bus/usb -v /home/developer/rknn-toolkit2/model_convert:/test rknn-toolkit2:2.3.2-cp38 /bin/bash

执行成功如下图所示:

4.3 模型转换操作说明

4.3.1 模型转换Demo目录结构

模型转换Demo组成如下图所示:

4.3.2 生成量化图片列表

在docker环境切换到模型转换工作目录:

cd /test/

如下图所示:

执行gen_list.py生成量化图片列表:

python gen_datas.py

命令行现象如下图所示:

生成“量化图片列表”如下文件夹所示:

4.3.3 onnx模型转换为rknn模型

onnx2rknn.py脚本默认do_quant = False不对模型进行量化，如果需要进行int8量化将do_quant = True，脚本代码清单如下所示：

from rknn.api import RKNN if __name__ == '__main__': h, w = 360, 640 do_quant = True quant = "fp16" if do_quant: quant = "int8" platform = "rv1126b" model_path = r"weight/xlsr_2x_{}x{}.onnx".format(h, w) output_path = r"weight/xlsr_2x_{}x{}-{}.rknn".format(h, w, quant) data_set = r'datas_{}x{}.txt'.format(h, w) # Create RKNN object rknn = RKNN(verbose=False) # Pre-process config print('--> Config model') rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform, optimization_level = 3, model_pruning = True, compress_weight = True, #sparse_infer = True, output_optimize = True, custom_string='output_format=nhwc', ) print('done') # Load model print('--> Loading model') ret = rknn.load_onnx(model=model_path, inputs=['image'], input_size_list=[[1, 3, h, w]]) if ret != 0: print('Load model failed!') exit(ret) print('done') # Build model print('--> Building model') ret = rknn.build(do_quantization=do_quant, dataset=data_set) if ret != 0: print('Build model failed!') exit(ret) print('done') # Export rknn model print('--> Export rknn model') ret = rknn.export_rknn(output_path) if ret != 0: print('Export rknn model failed!') exit(ret) print('done')

生成模型如下图所示，模型将保存至output_path 路径，此模型可以在rknn环境和EASY-EAI-Nano-TB环境运行：