Torch2Onnx

搞定动态输入!

上次保存的模型的作用是把输入图像分辨率给放大n倍这么一件事情,有一个问题就是放大的倍数是作为模型的初始化参数,在定义模型时传递进去的,一旦模型训练完成,放大倍数就无法调整了,这次博客的目的是实现动态缩放倍数这么一个功能。

  1. 很自然的想法
    修改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会报错
原因:torchonnx的过程中,所有的输入必须是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并且无法实现我们想要的功能。

  1. 怎么办?
    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)



你可能感兴趣的:(模型部署,深度学习,pytorch,python)