Torch2Onnx
搞定动态输入！
上次保存的模型的作用是把输入图像分辨率给放大n倍这么一件事情，有一个问题就是放大的倍数是作为模型的初始化参数，在定义模型时传递进去的，一旦模型训练完成，放大倍数就无法调整了，这次博客的目的是实现动态缩放倍数这么一个功能。
- 很自然的想法
修改torch模型，让缩放倍数作为一个正向传播时才确定的参数。
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并且无法实现我们想要的功能。
- 怎么办？
open-mmlab博客中给的方法是在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)