NAFNet网络图像去模糊及模型转为onnx
- NAFNet: Nonlinear Activation Free Network for Image Restoration（官网名称）
- 前言
- 一、NAFNet
- 二、使用步骤
- 1.demo模型验证
- 2.模型导出ONNX（重点官方没有）
- 总结
- 参考资料
前言
面朝大海，春暖花开
一、NAFNet
官方git:https://github.com/megvii-research/NAFNet/#results-and-pre-trained-models
官方介绍：
尽管最近在图像恢复领域取得了重大进展，但现有技术（SOTA）方法的系统复杂性也在增加，这可能会阻碍方法的方便分析和比较。在本文中，我们提出了一个简单的基线，它超越了 SOTA 方法并且计算效率很高。为了进一步简化基线，我们揭示了非线性激活函数，例如 Sigmoid、ReLU、GELU、Softmax 等不是必需的：它们可以被乘法替换或删除。因此，我们从基线推导出了一个非线性无激活网络，即 NAFNet。SOTA 结果是在各种具有挑战性的基准上实现的，例如 GoPro 上的 33.69 dB PSNR（用于图像去模糊），超过了之前的 SOTA 0.38 dB，计算成本仅为 8.4%；SIDD 上的 40.30 dB PSNR（用于图像去噪），超过之前的 SOTA 0.28 dB，计算成本不到一半。
按照官方的使用方法就可以。下面展示官方的一个demo代码
1.demo模型验证# ------------------------------------------------------------------------
# Copyright (c) 2022 megvii-model. All Rights Reserved.
# ------------------------------------------------------------------------
# Modified from BasicSR (https://github.com/xinntao/BasicSR)
# Copyright 2018-2020 BasicSR Authors
# ------------------------------------------------------------------------
import torch
# from basicsr.data import create_dataloader, create_dataset
from basicsr.models import create_model
from basicsr.train import parse_options
from basicsr.utils import FileClient, imfrombytes, img2tensor, padding, tensor2img, imwrite
# from basicsr.utils import (get_env_info, get_root_logger, get_time_str,
# make_exp_dirs)
# from basicsr.utils.options import dict2str
def main():
# parse options, set distributed setting, set ramdom seed
opt = parse_options(is_train=False)
opt['num_gpu'] = torch.cuda.device_count()
img_path = opt['img_path'].get('input_img')
output_path = opt['img_path'].get('output_img')
## 1. read image
file_client = FileClient('disk')
img_bytes = file_client.get(img_path, None)
try:
img = imfrombytes(img_bytes, float32=True)
except:
raise Exception("path {} not working".format(img_path))
img = img2tensor(img, bgr2rgb=True, float32=True)
## 2. run inference
opt['dist'] = False
model = create_model(opt)
model.feed_data(data={'lq': img.unsqueeze(dim=0)})
if model.opt['val'].get('grids', False):
model.grids()
model.test()
if model.opt['val'].get('grids', False):
model.grids_inverse()
visuals = model.get_current_visuals()
sr_img = tensor2img([visuals['result']])
imwrite(sr_img, output_path)
print(f'inference {img_path} .. finished. saved to {output_path}')
if __name__ == '__main__':
main()
单个图像推理：
python basicsr/demo.py -opt options/test/REDS/NAFNet-width64.yml --input_path ./demo/blurry.jpg --output_path ./demo/deblur_img.png2.模型导出ONNX（重点官方没有）
在转成onnx需要更改模型网络的一些代码：NAFNet-mainbasicsrmodelsarchsNAFNet_arch.py
需要更改的用##########包围，供大家参考。
class NAFBlock(nn.Module):
def __init__(self, c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.):
super().__init__()
dw_channel = c * DW_Expand
self.conv1 = nn.Conv2d(in_channels=c, out_channels=dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
self.conv2 = nn.Conv2d(in_channels=dw_channel, out_channels=dw_channel, kernel_size=3, padding=1, stride=1, groups=dw_channel,
bias=True)
self.conv3 = nn.Conv2d(in_channels=dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
# Simplified Channel Attention
self.sca = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels=dw_channel // 2, out_channels=dw_channel // 2, kernel_size=1, padding=0, stride=1,
groups=1, bias=True),
)
# SimpleGate
self.sg = SimpleGate()
ffn_channel = FFN_Expand * c
self.conv4 = nn.Conv2d(in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
self.conv5 = nn.Conv2d(in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True)
# self.norm1 = LayerNorm2d(c)
# self.norm2 = LayerNorm2d(c)
###########################
self.norm1 = torch.nn.LayerNorm(c)
self.norm2 = torch.nn.LayerNorm(c)
###########################
self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
def forward(self, inp):
x = inp
##########################
x = torch.permute(x, (0, 3, 2, 1))
############################
x = self.norm1(x)
##########################
x = torch.permute(x, (0, 3, 2, 1))
#############################
x = self.conv1(x)
x = self.conv2(x)
x = self.sg(x)
x = x * self.sca(x)
x = self.conv3(x)
x = self.dropout1(x)
y = inp + x * self.beta
################################
yy = torch.permute(y, (0, 3, 2, 1))
yy = self.norm2(yy)
x = torch.permute(yy, (0, 3, 2, 1))
x = self.conv4(x)
####################################
# x = self.conv4(self.norm2(y))
x = self.sg(x)
x = self.conv5(x)
x = self.dropout2(x)
return y + x * self.gamma
代码如下（示例）：
import os
import torch
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from basicsr.models import create_model
from basicsr.train import parse_options
from basicsr.utils import FileClient, imfrombytes, img2tensor, padding, tensor2img, imwrite
from basicsr.models.archs import define_network
from copy import deepcopy
from torch.nn.parallel import DataParallel, DistributedDataParallel
def model_to_device(net):
"""Model to device. It also warps models with DistributedDataParallel
or DataParallel.
Args:
net (nn.Module)
"""
opt = parse_options(is_train=False)
opt['num_gpu'] = torch.cuda.device_count()
device = torch.device('cuda' if opt['num_gpu'] != 0 else 'cpu')
net = net.to(device)
if opt['dist']:
find_unused_parameters = opt.get('find_unused_parameters',
False)
net = DistributedDataParallel(
net,
device_ids=[torch.cuda.current_device()],
find_unused_parameters=find_unused_parameters)
elif opt['num_gpu'] > 1:
net = DataParallel(net)
return net
def main():
# parse options, set distributed setting, set ramdom seed
opt = parse_options(is_train=False)
opt['num_gpu'] = torch.cuda.device_count()
img_path = opt['img_path'].get('input_img')
output_path = opt['img_path'].get('output_img')
## 1. read image
file_client = FileClient('disk')
img_bytes = file_client.get(img_path, None)
try:
img = imfrombytes(img_bytes, float32=True)
except:
raise Exception("path {} not working".format(img_path))
img = img2tensor(img, bgr2rgb=True, float32=True)
## 2. run inference
opt['dist'] = False
model = create_model(opt)
# 把模型使用的网络单独拿出来转为onnx
net_g = define_network(deepcopy(opt['network_g']))
net_g = model_to_device(net_g)
dummy_input = torch.randn(1, 3, 280, 280, device="cuda")
input_names = ["actual_input_1"] + ["learned_%d" % i for i in range(16)]
output_names = ["output1"]
output_onnx_path = "E:\NAFNet-main\NAFNet.onnx"
torch.onnx.export(net_g, dummy_input, output_onnx_path,
verbose=True,input_names=input_names, output_names=output_names, opset_version=11)
model.feed_data(data={'lq': img.unsqueeze(dim=0)})
if model.opt['val'].get('grids', False):
model.grids()
model.test()
if model.opt['val'].get('grids', False):
model.grids_inverse()
visuals = model.get_current_visuals()
sr_img = tensor2img([visuals['result']])
imwrite(sr_img, output_path)
print(f'inference {img_path} .. finished. saved to {output_path}')
if __name__ == '__main__':
main()
同样执行上面demo运行方法，就可以将模型转为onnx了
python basicsr/demo.py -opt options/test/REDS/NAFNet-width64.yml --input_path ./demo/blurry.jpg --output_path ./demo/deblur_img.png
导出后模型进行验证，这里注意dummy_input = torch.randn(1, 3, 280, 280, device=“cuda”)你打包时候设置输入图像的大小（280， 280），做预测就要符合图像大小要求。
验证onnx的代码如下：
import onnxruntime
from onnxruntime.datasets import get_example
import torch
from basicsr.utils import FileClient, imfrombytes, img2tensor, padding, tensor2img, imwrite
# 测试数据
# dummy_input = torch.randn(1, 3, 256, 256, device='cpu')
img_path = r"E:NAFNet-maindatasetscaisitraininput1-35.jpg"
output_path = r"E:oilworkNAFNet-mainNAFNet-main"
file_client = FileClient('disk')
img_bytes = file_client.get(img_path, None)
img = imfrombytes(img_bytes, float32=True)
img = img2tensor(img, bgr2rgb=True, float32=True).unsqueeze(dim=0)
print(img.shape)
example_model = get_example(r"E:oilworkNAFNet-mainNAFNet.onnx")
# netron.start(example_model) 使用 netron python 包可视化网络
sess = onnxruntime.InferenceSession(example_model, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
imgs =img.cpu().numpy()
# print(imgs)
# onnx 网络输出
input_names = ["actual_input_1"] + ["learned_%d" % i for i in range(16)]
onnx_out = sess.run(None, {input_names[0]: imgs})
print(onnx_out)
这里的结果只是打印输出的矩阵，自己需要按照解码方式进行转换。
总结
模型转为onnx的时候，需要将模型的网络单独拿出来，最好是和数据处理分开，这样模型在转为onnx出现问题比较少。
参考资料https://github.com/megvii-research/NAFNet/#results-and-pre-trained-models
