一、背景
最近做的一个项目需求，数据上的一个特点是需要检测的目标很小，且部分情况下由于画面不清晰导致目标与背景部分比较难区分，需要想办法提高目标区域的清晰度，加大目标特征在画面中的占比，于是想到了图像超分重建，首先粗提取目标区域，再使用超分放大目标区域。
目前常见的一些超分模型的重建过程都相对较慢，如果想在实时视频流的推理过程中加入超分模型，还能不能满足实时性要求？或者说对帧率影响多大？如果再将其移植到算力更低的边端设备，效果又会如何？这就需要把模型运行起来并一一验证了。
由于采集到的数据经过提取后都是低分辨率图像，所以就需要把预训练好的模型拿过来直接使用。如果想要在自己的数据集上重新训练模型，或者是微调，高分辨率图像是必不可少的，这与实际情况相悖（如果有高分辨率图像那也用不着重建了）。综合以上情况最终选择的是 Real-ESRGAN。本文内容不含模型训练和微调部分，只记录使用官方模型进行重建和模型量化后的对比测试。
二、本地测试
Real-ESRGAN 是由腾讯 ARC 实验室发布的一个盲图像超分辨率模型，模拟图像从高分辨率到低分辩率过程中的各种退化，然后通过低清图倒推出它的高清图 [1]。Real-ESRGAN 是 ESRGAN 针对真实世界图像重建优化的升级版本，重建效果更好更稳定，二者都是基于 BasicSR （一个基于 PyTorch 的开源图像视频复原工具箱）进行训练和推理的。
仓库地址：https://github.com/xinntao/Real-ESRGAN
1. 环境安装
# BasicSR 用于训练和推理
pip install basicsr
# facexlib and gfpgan 用于面部增强
pip install facexlib
pip install gfpgan
pip install -r requirements.txt
python setup.py develop2. 预训练模型
截止到目前为止官方提供的几类模型：
（1）适用一般图像
| 模型 | 对应判别器（用于微调） | 描述 |
| RealESRGAN_x4plus | RealESRGAN_x4plus_netD | 4倍模型，默认使用 |
| RealESRGAN_x2plus | RealESRGAN_x2plus_netD | 2倍模型 |
| RealESRNet_x4plus | 无 | 使用 MSE Loss 训练的4倍模型 |
| official ESRGAN_x4 | 无 | 4倍 ESRGAN 模型 |
| realesr-general-x4v3 | 无 | 轻量化4倍模型，可调节降噪程度（最新） |
| realesr-general-wdn-x4v3 | 无 | 配合 realesr-general-x4v3 使用的降噪模型 |
虽然官方说 realesr-general-x4v3 也支持1-3倍，但实际这个1-3倍是对重建后的图像做resize，实际模型直接输出的还是只有4倍
（2）适用动漫图像
| 模型 | 对应判别器（用于微调） | 描述 |
| RealESRGAN_x4plus_anime_6B | RealESRGAN_x4plus_anime_6B_netD | 针对动漫图像优化的4倍模型 |
（3）适用动漫视频
| 模型 | 对应判别器（用于微调） | 描述 |
| realesr-animevideov3 | 无 | 1-4倍模型 |
下载之后放到 weights 文件夹下。本文只测试了 RealESRGAN_x4plus 和 realesr-general-x4v3，测试图像是真实世界图像，测试效果前者更好，后者推理速度更快，效果也还可以，后续量化使用的也是 realesr-general-x4v3。
3. 测试方式
官方提供三种测试方式：
（1）在线 Demo
- ARC Demo（仅支持动漫图，没试）
- Colab Demo1（支持一般图像，试了但环境问题没跑通）
- Colab Demo2（针对动漫视频，没试）
（2）可执行文件
支持 Win/Linux/Mac，可直接安装使用，没试。
（3）python 脚本执行
本文主要使用这种方式。
执行命令，输出结果保存到 results 文件夹
# -n 模型名称；-i 测试数据所在文件夹
python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs4. 结果对比
（1）效果对比
比较了 RealESRGAN_x4plus、realesr-general-x4v3（去噪水平 0.5） 和 realesr-general-x4v3 （以下简称 x4plus、x4v3 和 x4v3-onnx）导出为静态 onnx 的重建效果。
重建效果 x4plus > x4v3 > x4v3-onnx，如果是动态 onnx 模型，效果几乎无损（肉眼观察）。
（2）速度对比
由于 x4plus 和 x4v3 是动态输入，推理速度随图像分辨率变化，所以仅以上面的两张图为例进行测试
| 模型 | 图1 | 图2 |
| x4plus | 0.88s | 2.31s |
| x4v3 | 0.08s | 0.31s |
| x4v3-onnx | 0.51s | 0.65s |
可以看到 x4v3 的确比 x4plus 快不少，重建效果也还不错，而 x4v3-onnx 使用的是 CPU 推理，速度方面逊于受 GPU 加持的 x4v3。
x4v3 训练默认使用 prelu 激活函数，将其改为 relu 并导出 onnx，推理速度没有提升但效果明显变差。
三、量化部署
1. pt --> onnx
scripts/pytorch2onnx.py 为官方提供的转换脚本，但只写了 x4plus 的转换，现在需要转换 x4v3，所以要修改一下脚本。
import argparse
import torch
import torch.onnx
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
def dni(net_a, net_b, dni_weight, key='params', loc='cpu'):
"""Deep network interpolation.
``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition``
"""
net_a = torch.load(net_a, map_location=torch.device(loc))
net_b = torch.load(net_b, map_location=torch.device(loc))
for k, v_a in net_a[key].items():
net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k]
return net_a
def main(args):
# x4v3使用SRVGGNetCompact
if 'v3' in args.input:
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
keyname = 'params'
dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
pth = dni(args.input, args.input2, dni_weight)
else:
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
keyname = 'params_ema'
pth = torch.load(args.input)
state_dict = pth[keyname]
model.load_state_dict(state_dict)
model.train(False)
model.cpu().eval()
# 固定输入大小，导出为静态模型
x = torch.rand(1, 3, 256, 256)
# 导出
with torch.no_grad():
torch.onnx.export(model,
x,
args.output,
opset_version=11,
)
print('Export ONNX model successful!')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--input', type=str, default='', help='Input model path')
parser.add_argument(
'--input2', type=str, default='', help='Input denoise model path')
parser.add_argument('--output', type=str, default='', help='Output onnx path')
parser.add_argument(
'-dn',
'--denoise_strength',
type=float,
default=0.5,
help=('Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. '
'Only used for the realesr-general-x4v3 model'))
args = parser.parse_args()
main(args)
如果导出动态模型，需要在导出前添加动态参数，这样导出的模型可支持任意宽高图像输入。
x = torch.rand(1, 3, 256, 256)
dynamic = {'images': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'},
'output': {0: 'batch_size', 1: 'channel', 2: 'height', 3: 'width'}}
# Export the model
with torch.no_grad():
torch.onnx.export(model,
x,
args.output,
opset_version=11,
input_names=['images'],
output_names=['output'],
dynamic_axes=dynamic,
)如果把 prelu 改为 relu，需要修改 act_type 和权重文件（还是不要改了，效果不好 - -）
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='relu')state_dict = pth[keyname]
# 需要删除部分层
state_dict_copy = state_dict.copy()
del_key = [str(i) for i in range(1, 66, 2)]
for k, v in state_dict.items():
k2 = k.split('.')[1]
k3 = k.split('.')[2]
if k3 == 'weight':
if k2 in del_key:
del state_dict_copy[k]
model.load_state_dict(state_dict_copy)2. onnx --> rknn
一开始是想把 onnx 模型直接放到盒子上，使用 onnxruntime 推理，这样没法调用 npu 加速，实测推理速度在1.7s，有点惨不忍睹，所以要用起来就得想办法转成rknn。RK 官方的 Model Zoo 里没有 Real-ESRGAN，只能尝试一下看能不能转换成功，没想到使用现有脚本（可参考RK官方脚本或看我之前的文章）直接就能转，也可以直接做量化。
下面记录一下对 x4v3 做转换和量化后的模型（分别记作 x4v3-rknn 和 x4v3-rknn-q）测试重建效果和推理速度。
（1）重建效果
可以看到没有量化的 x4v3-rknn 与 x4v3-onnx 的重建效果没有肉眼可见的差别，但量化后的 x4v3-rknn-q 对第一张图像的重建效果与 x4v3-rknn 相比就比较明显，背景没有 x4v3-rknn 平滑，将降噪水平调到最高，也没有改善很多，可能是量化对模型精度的损失较大，但第二张图的效果差别又不明显了。
（2）推理速度
以下测试均在 RK3588 上进行。
| 模型 | 图一 | 图二 |
| x4v3-onnx | 1.71s | 1.70s |
| x4v3-rknn | 0.14s (91.8%↑) | 0.15s (91.1%↑) |
| x4v3-rknn-q | 0.10s (94.1%↑) | 0.11s (93.5%↑) |
经过 NPU 加速后推理速度提升很明显。
四、总结
经过一番折腾实现了 Real-ESRGAN 的边端推理，过程比预想的顺利，最终的效果和速度（能力有限，目前只能做到这个速度）基本满足需求。除了 Real-ESRGAN，也测试了 EGVSR 和CAMixerSR，都是主打速度快。
在 GPU 上对同样 13 张不同分辨率的图像做4倍重建，x4plus 总耗时 16s，x4v3 总耗时 4.59s，CMixerSRx4 总耗时 12.68s，TecoGAN_BD_iter50000 总耗时 5.23s，效果上 CMixerSRx4 还是比 x4v3 差一些，EGVSR 可能是我用法不对，重建出来的图像没法看，后面有时间再搞搞清楚。
参考资料
[1] 盲图像超分辨率模型 Real-ESRGAN 使用教程-知乎
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
