OpenVINO 2021r2 C++ 超分辨率重建EDSR

最新推荐文章于 2024-04-02 19:20:14 发布

原创

最新推荐文章于 2024-04-02 19:20:14 发布 · 973 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#超分辨率重建 #OpenVINO

本文介绍如何将EDSR超分辨率模型从PyTorch转换为ONNX格式，并使用OpenVINO进行部署。涵盖模型转换步骤、预处理及后处理方法、性能对比等内容。

接下来试试EDSR，NTIRE2017 超分辨率挑战赛的第一名

随便从网上找了个模型 https://github.com/achie27/super-resolution ，在这个项目里EDSR是用pytorch实现的，对于OpenVINO来说，MO不支持pth格式的模型，必须转换到ONNX格式的模型才行。所以要利用这个项目里自带的to_onnx.py来做。

首先对原项目里的to_onnx.py做一下修改，原始项目里的代码是基于有cuda加速的pytorch写的，对于我这台只有集成显卡的老破小来说需要改成CPU的版本.

import torch
from architecture import arch

settings = {
	"model"                 : "esrgan", #srcnn, fsrcnn, espcn, edsr, srgan, esrgan, or prosr
	"edsr"                 : 2             # 2 or 4
}

def main():
	model = arch(settings['model'], settings['scale'], False).getModel()
	model.eval()

	dummy_input = torch.randn(1, 3, 480, 640, device = 'cpu')
	output_names = ["the_output"] # will use this in DnnSuperResolution

	with torch.no_grad():
		torch.onnx.export(
			model, 
			dummy_input, 
			settings['model'] + '_x' + str(settings['scale']) + '.onnx',
			output_names = output_names, 
			verbose=True
		)

if __name__ == '__main__':
    main()

运行

python to_onnx.py

得到edsr_x2.onnx

接下来看看edsr推理对于数据的预处理和后处理部分, 在原始项目的architecture\edsr\__init__.py里，

class MeanShift(nn.Conv2d):
    def __init__(
        self, rgb_range,
        rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):

        super(MeanShift, self).__init__(3, 3, kernel_size=1)
        std = torch.Tensor(rgb_std)
        self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)
        self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) / std
        for p

最低0.47元/天解锁文章