上次保存的模型的作用是把输入图像分辨率给放大n倍这么一件事情,有一个问题就是放大的倍数是作为模型的初始化参数,在定义模型时传递进去的,一旦模型训练完成,放大倍数就无法调整了,这次博客的目的是实现动态缩放倍数这么一个功能。
class SuperResolutionNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
self.relu = nn.ReLU()
def forward(self, x, upscale_factor):
x = interpolate(x,
scale_factor=upscale_factor,
mode='bicubic',
align_corners=False)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.conv3(out)
return out
output_torch = model(x, 3).detach().numpy()
torch.onnx.export(model,
(x, 3)
'srcnn_dynamic.onnx',
opset_version = 11,
input_names = ['input', 'factor'],
output_names = ['output'])
发现torch
能够正常输出,但是onnx
会报错
原因:torch
转onnx
的过程中,所有的输入必须是tensor
类型,这是规定,所以下面把model
的输入变成tensor
试一下
x = F.interpolate(x, scale_factor=upscale_factor.item(), mode='bicubic', align_corners=False)
发现虽然.onnx
模型导出来了,报了条Warning
,并且通过Netron
查看网络的输入还是只有一个input
,并没有缩放倍数upscale_factor
。
原因:在模型的forward
过程中将执行了tensor.item()
这个操作,onnx
无法记录这一操作所以会导致Warning
并且无法实现我们想要的功能。
torch --> onnx
这条技术栈中,onnx
对于pytorch
中所有的插值操作都是由Resize
这个算子实现的。换句话说所有pytorch
中的插值函数都会经过一个映射函数,将pytorch
中的插值函数映射成onnx
中的Resize
函数。问题就出在这里目前的pytorch --> onnx
的映射过程中,映射成Resize
算子中的scale
参数为常数,无法作为一个变量。onnx
中的Resize
算子上,这样onnx
中新的Resize
算子就可以支持动态缩放了。"""
用pytorch搭建一个深度神经网络(动态缩放倍数),转成onnx格式,用onnxruntime进行推理
分析:
动态插值算子在pytorch中没有实现
插值算子在onnx中的形式为Resize
需要自己在pytorch中定义一个可以携带放大参数的插值算子,然后写一个映射函数,映射到onnx中存在的Resize函数即可
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import onnxruntime
import cv2
import numpy as np
"""
定义一个pytorch中的动态插值算子
定义一个pytorch到onnx的映射关系
"""
class NewInterpolate(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.tensor, scales: torch.tensor): # 该算子在pytorch中的推理函数定义,第一个参数必须是ctx
scales = scales.tolist()[-2:] # 将[1, 1, scales, scales]转换成[scales, scales]
return F.interpolate(input,
scale_factor = scales,
mode = 'bicubic',
align_corners = False)
@staticmethod
def symbolic(g, input, scales): # 该算子的映射函数,symbolic方法第一个参数必须是g,之后的参数是算子的自定义输入
return g.op(
# NODE PROPERTIES
"Resize",
# INPUTS
input,
g.op("Constant", value_t = torch.tensor([], dtype = torch.float32)),
scales,
# ATTRIBUTES
coordinate_transformation_mode_s = 'pytorch_half_pixel',
mode_s = 'cubic',
nearest_mode_s = 'floor'
)
class SuperResolutionNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
self.relu = nn.ReLU()
def forward(self, x: torch.tensor, upscale_factor: torch.tensor):
x = NewInterpolate.apply(x, upscale_factor)
# x = F.interpolate(x, scale_factor=upscale_factor.item(), mode='bicubic', align_corners=False)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.conv3(out)
return out
if __name__ == '__main__':
model = SuperResolutionNet()
model.eval()
x = torch.rand(1, 3, 256, 256)
output_torch = model(x, torch.tensor([1, 1, 3, 3], dtype=torch.float)).detach().numpy()
print(output_torch.shape)
# 把model转换成onnx格式并保存
with torch.no_grad():
torch.onnx.export(model,
(x, torch.tensor([1, 1, 3, 3], dtype=torch.float)), # 注意这里的输入数据的类型必须全部为tensor
'srcnn_dynamic.onnx',
opset_version = 11,
input_names = ['input', 'factor'],
output_names = ['output'])
session = onnxruntime.InferenceSession('srcnn_dynamic.onnx')
inputs_ort = {'input': x.numpy(), 'factor': np.array([1, 1, 3, 3], dtype=np.float32)} # 注意onnxruntime的输入的类型必须是numpy格式,不能是tensor
output_ort = session.run(['output'], inputs_ort)[0]
print(output_ort.shape)