Pytorch模型部署：torch-onnx-ncnn

主要是学习记录以及一些理解，大部分内容来自以下链接，如有侵权请联系我

学习内容：

OpenMMlab所著《模型部署那些事》专栏，推荐一看！！！
TorchScript简介
TorchScript与onnx的区别与联系

模型部署简介
将一个torch模型转为onnx并与原模型的输出进行对比

解决模型部署中的难题
让模型的输入动态化，将固定缩放倍数的模型修改为可以动态指定倍数

Pytorch转onnx
torch.onnx.export函数详解及实战

算子支持
将ATen、TorchScript算子与onnx算子进行映射
添加C++拓展并与onnx算子进行映射，使用torch.autograd.Function进行封装，在torch中使用该自定义算子，再转为onnx

onnx读写与调试
onnx读写、修改

环境：

Ubuntu 20.04.4 LTS
Docker Client Version—20.10.16
Python — 3.7.13
GCC — 7.5.0
torch — 1.8.0+cpu
onnx — 1.11.0
protobuf — 3.20.1
ncnn — 1.0.20220519
编辑器VScode，安装的拓展：Docker(管理镜像及容器)、Remote-Containers(连接远程仓库)、C/C++

平时我们主要接触的是深度学习框架如pytorch等，在这些框架上训练模型，通过反复迭代得到最终的模型确定各个节点参数，之后将训练完毕的模型转为中间表示，如onnx、caffe，针对网络结构的优化（蒸馏、剪枝、量化等）会在中间表示上进行；再使用推理引擎将中间表示转化为硬件平台所需格式。
常见的推理引擎对应的硬件平台：

• ONNXRuntime: ONNX Runtime是一个跨平台的机器学习训练推理加速器，通过图形优化和变换以及硬件加速器提供优秀的推理性能。拥有完善的对ONNX的支持。ONNX Runtime 是直接对接 ONNX 的，即 ONNX Runtime 可以直接读取并运行 .onnx 文件, 而不需要再把 .onnx 格式的文件转换成其他格式的文件。也就是说，对于 PyTorch - ONNX - ONNX Runtime 这条部署流水线，只要在目标设备中得到 .onnx 文件，并在 ONNX Runtime 上运行模型，模型部署就算大功告成了。
• TensorRT: NVIDIA® TensorRT™ 是一个用于高性能深度学习推理的开发工具包（SDK）。借助Nvidia的设备特性，TensorRT可以优化模型的推理，提供更低的推理延迟以及更高的吞吐量。如果您希望将模型部署在NVIDIA硬件设备上，那么TensorRT就是一个合适的选择。
• ncnn: ncnn 是一个为手机端极致优化的高性能神经网络前向计算框架。ncnn 从设计之初深刻考虑手机端的部署和使用。无第三方依赖，跨平台。基于 ncnn，开发者能够将深度学习算法轻松移植到手机端高效执行，开发出人工智能 APP，将 AI 带到您的指尖。
• PPLNN: PPLNN是一个为高效AI推理所开发的高性能深度学习推理引擎。可以用于各种ONNX模型的推理。并且对OpenMMLab有非常强的支持。
• OpenVINO: OpenVINO™ 是一个为优化与部署AI推理开发的开源工具集。该工具集可无缝集成到 Intel 硬件平台，包括最新的神经网络加速芯片，Intel计算棒，边缘设备等。

受硬件条件限制，我打算将模型部署在安卓端，部署路线主要如下所示

环境配置

Docker安装及镜像仓库地址配置

然后拉取我使用的是MMedploy中docker/CPU文件夹下的dockerfile在容器中进行配置，官方的安装教程点这里

首先clone仓库，如果克隆失败建议切换到国内源，然后编译 MMDeploy（在这里卡了很长时间）主要是各种包装不上以及拒绝链接等问题，我的解决方法见这里和这里，这个dockerfile中安装了很多需要的命令如curl，以及torch、ncnn等接下来必备的包，同时完成了编译，非常方便

个人先安装官方的教程捋了一遍，没有使用自己的模型

总结及感悟

TorchScript简介部分

这张图很清楚的说明了onnx与torchscript的关系，这两个都是作为中间表示，trace和script两种方法都可以进行序列化，序列化后的模型不再与 python 相关，可以被部署到各种平台上。torch.onnx.export函数可以把 PyTorch 模型转换成 ONNX 模型，这个函数会使用 trace 的方式记录 PyTorch 的推理过程。

onnx的生成主要包括三步：

1. 使用 trace 的方式先生成一个 TorchScipt 模型，如果你转换的本身就是 TorchScript 模型，则可以跳过这一步。
2. 使用许多 pass 对 1 中生成的模型进行变换，其中对 ONNX 导出最重要的一个 pass 就是ToONNX，这个 pass会进行一个映射，将 TorchScript 中prim、aten空间下的算子映射到onnx空间下的算子。
3. 使用 ONNX 的 proto格式对模型进行序列化，完成 ONNX 的导出。

也就是说onnx是一种中间表示，而torchscript相当于是torch模型转为onnx第一步的结果。

模型部署简介部分

pytorch模型转为onnx

import os 
 
import cv2 
import numpy as np 
import requests 
import torch 
import torch.onnx 
from torch import nn 
 
class SuperResolutionNet(nn.Module): 
    def __init__(self, upscale_factor): 
        super().__init__() 
        self.upscale_factor = upscale_factor 
        self.img_upsampler = nn.Upsample( 
            scale_factor=self.upscale_factor, 
            mode='bicubic', 
            align_corners=False) 
 
        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): 
        x = self.img_upsampler(x) 
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out 
 
# Download checkpoint and test image 
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 
    'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png'] 
names = ['srcnn.pth', 'face.png'] 
for url, name in zip(urls, names): 
    if not os.path.exists(name): 
        open(name, 'wb').write(requests.get(url).content) 
 
def init_torch_model(): 
    torch_model = SuperResolutionNet(upscale_factor=3) 
 
    state_dict = torch.load('srcnn.pth')['state_dict'] 
 
    # Adapt the checkpoint 
    for old_key in list(state_dict.keys()): 
        new_key = '.'.join(old_key.split('.')[1:]) 
        state_dict[new_key] = state_dict.pop(old_key) 
 
    torch_model.load_state_dict(state_dict) 
    torch_model.eval() 
    return torch_model 
 
model = init_torch_model() 
input_img = cv2.imread('face.png').astype(np.float32) 
 
# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0) 
 
# Inference 
torch_output = model(torch.from_numpy(input_img)).detach().numpy() 
 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
 
# Show image 
cv2.imwrite("face_torch.png", torch_output)

#转为onnx
x = torch.randn(1, 3, 256, 256) 
 
with torch.no_grad(): 
    torch.onnx.export( 
        model, 
        x, 
        "srcnn.onnx",  #导出onnx名
        opset_version=11,  #算子集版本
        input_names=['input'], 
        output_names=['output'])

导出的onnx可以使用netron软件来查看
onnx中每个算子记录了算子属性、图结构、权重三类信息。

对于卷积来说，算子属性包括了卷积核大小(kernel_shape)、卷积步长(strides)等内容。这些算子属性最终会用来生成一个具体的算子。

图结构信息指算子节点在计算图中的名称、邻边的信息。对于图中的卷积来说，该算子节点叫做 Conv_2，输入数据叫做 11，输出数据叫做 12。根据每个算子节点的图结构信息，就能完整地复原出网络的计算图。

权重信息指的是网络经过训练后，算子存储的权重信息。对于卷积来说，权重信息包括卷积核的权重值和卷积后的偏差值。点击图中 conv1.weight, conv1.bias 后面的加号即可看到权重信息的具体内容。

使用onnxruntime部署

import onnxruntime 
 
ort_session = onnxruntime.InferenceSession("srcnn.onnx")  #获取一个 ONNX Runtime 推理器
ort_inputs = {'input': input_img} 
ort_output = ort_session.run(['output'], ort_inputs)[0] 
'''
推理器的 run 方法用于模型推理，其第一个参数为输出张量名的列表，第二个参数为输入值的字典。
其中输入值字典的 key 为张量名，value 为 numpy 类型的张量值。
输入输出张量的名称需要和torch.onnx.export 中设置的输入输出名对应。
'''
ort_output = np.squeeze(ort_output, 0) 
ort_output = np.clip(ort_output, 0, 255) 
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8) 
cv2.imwrite("face_ort.png", ort_output)

解决模型部署中的难题部分

这一节主要是讲了如何实现模型的动态化输入，将之前的srcnn的缩放倍数也作为转onnx的输入

在 PyTorch 中无论是使用最早的 nn.Upsample，还是后来的 interpolate，PyTorch 里的插值操作最后都会转换成 ONNX 定义的 Resize 操作。也就是说，所谓 PyTorch 转 ONNX，实际上就是把每个 PyTorch 的操作映射成了 ONNX 定义的算子。
但是Resize 算子的 scales 只能是常量，无法满足我们的需求。我们可以自己定义一个实现插值的 PyTorch 算子，然后让它映射到一个我们期望的 ONNX Resize 算子上。将他重新封装成一个新算子，实质上还是使用resize算子，不过将scales改为一个变量

完整代码如下：

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
 
 #定义新算子，使用torch.autograd.Function封装
class NewInterpolate(torch.autograd.Function): 
 
	#映射到 ONNX 的方法由一个算子的 symbolic 方法决定。symbolic 方法第一个参数必须是g。
	#ONNX 算子的具体定义由 g.op 实现。g.op 的每个参数都可以映射到 ONNX 中的算子属性（算子记录的三个信息之一）
    @staticmethod 
    def symbolic(g, input, scales): #此处Resize为onnx中算子名称
        return g.op("Resize", 
                    input, 
                    g.op("Constant", 
                         value_t=torch.tensor([], dtype=torch.float32)), 
                    scales, 
                    coordinate_transformation_mode_s="pytorch_half_pixel", 
                    cubic_coeff_a_f=-0.75, 
                    mode_s='cubic', 
                    nearest_mode_s="floor") 
                    
 	#推理行为由算子的 foward 方法决定。该方法的第一个参数必须为 ctx
    @staticmethod 
    def forward(ctx, input, scales): 
        scales = scales.tolist()[-2:] 
        return interpolate(input, 
                           scale_factor=scales, 
                           mode='bicubic', 
                           align_corners=False) 
 
 
class StrangeSuperResolutionNet(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 = NewInterpolate.apply(x, upscale_factor) 
        #apply是torch.autograd.Function 的一个方法，这个方法完成了 Function 在前向推理或者反向传播时的调度。我们在使用 Function 的派生类做推理时，不应该显式地调用 forward()，而应该调用其 apply 方法。
        out = self.relu(self.conv1(x)) 
        out = self.relu(self.conv2(out)) 
        out = self.conv3(out) 
        return out 
 
 
def init_torch_model(): 
    torch_model = StrangeSuperResolutionNet() 
 
    state_dict = torch.load('srcnn.pth')['state_dict'] 
 
    # Adapt the checkpoint 
    for old_key in list(state_dict.keys()): 
        new_key = '.'.join(old_key.split('.')[1:]) 
        state_dict[new_key] = state_dict.pop(old_key) 
 
    torch_model.load_state_dict(state_dict) 
    torch_model.eval() 
    return torch_model 
 
 
model = init_torch_model() 
factor = torch.tensor([1, 1, 3, 3], dtype=torch.float) 
 
input_img = cv2.imread('face.png').astype(np.float32) 
 
# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0) 
 
# Inference 
torch_output = model(torch.from_numpy(input_img), factor).detach().numpy() 
 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
 
# Show image 
cv2.imwrite("face_torch_3.png", torch_output) 

#导出onnx
x = torch.randn(1, 3, 256, 256) 
 
with torch.no_grad(): 
    torch.onnx.export(model, (x, factor), 
                      "srcnn3.onnx", 
                      opset_version=11, 
                      input_names=['input', 'factor'], 
                      output_names=['output']) 


#使用onnxruntime部署
import onnxruntime 

#此时就可以在部署时再指定input_factor来修改缩放维度，而不需要在导出onnx之前修改
input_factor = np.array([1, 1, 4, 4], dtype=np.float32) 
ort_session = onnxruntime.InferenceSession("srcnn3.onnx") 
ort_inputs = {'input': input_img, 'factor': input_factor} 
ort_output = ort_session.run(None, ort_inputs)[0] 
 
ort_output = np.squeeze(ort_output, 0) 
ort_output = np.clip(ort_output, 0, 255) 
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8) 
cv2.imwrite("face_ort_3.png", ort_output) 

简单总结来说，想实现模型在转为onnx之后，部署时的动态化输入，可以定义新算子，使用forward函数定义推理行为，使用symbolic函数定义算子，将需要动态化输入的参数作为forward的参数之一，之后用torch.autograd.function封装。之后在pytorch模型中调用即可

Pytorch转onnx

这一节详细介绍了详细介绍 PyTorch 到 ONNX 的转换函数—— torch.onnx.export

torch.onnx.export中需要的模型实际上是一个torch.jit.ScriptModule。而要把普通 PyTorch 模型转一个这样的 TorchScript 模型，有跟踪（trace）和记录（script）两种导出计算图的方法。如果给torch.onnx.export传入了一个普通 PyTorch 模型（torch.nn.Module)，那么这个模型会默认使用跟踪的方法导出。
跟踪法只能通过实际运行一遍模型的方法导出模型的静态图，即无法识别出模型中的控制流（如循环）；记录法则能通过解析模型来正确记录所有的控制流。记录法导出模型不需要实际运行。两种方法在遇到循环时的处理方法不同，trace方法由于是通过跟踪实现，那么不同的循环次数得到的onnx结构不同，而记录法并不实际运行，以Loop节点表示循环，不同的循环次数，onnx结构都一致。

torch.onnx.export常见参数解释：

• input_names, output_names
  ONNX 模型的每个输入和输出张量都有一个名字。很多推理引擎在运行 ONNX 文件时，都需要以“名称-张量值”的数据对来输入数据，并根据输出张量的名称来获取输出数据。在进行跟张量有关的设置（比如添加动态维度）时，也需要知道张量的名字。在实际的部署流水线中，我们都需要设置输入和输出张量的名称，并保证 ONNX 和推理引擎中使用同一套名称。

• export_params
  模型中是否存储模型权重。一般中间表示包含两大类信息：模型结构和模型权重。ONNX 是用同一个文件表示记录模型的结构和权重的。部署时一般都默认这个参数为 True。如果 onnx 文件是用来在不同框架间传递模型（比如 PyTorch 到 Tensorflow）而不是用于部署，则可以令这个参数为 False。

• opset_version
  参考的onnx算子集版本

• dynamic_axes
  指定输入输出张量的哪些维度是动态的。
  为了追求效率，ONNX 默认所有参与运算的张量都是静态的（张量的形状不发生改变）。但在实际应用中，我们又希望模型的输入张量是动态的，尤其是本来就没有形状限制的全卷积模型。因此，我们需要显式地指明输入输出张量的哪几个维度的大小是可变的。注意：ONNX 要求每个动态维度都有一个名字
  例子：对于形状为(1, 3, 10, 10)的输入，如果想让他的第二维度为动态输入，则可以设置

dynamic_axes_0 = { 
    'in' : {2: 'height'}, 
    'out' : {2: 'height'} 
} 

将dynamic_axes_0作为torch.onnx.export的参数导出即可，此时只能是指定了动态维度的input可变，其他维度如果发生变化，会导出失败。比如我们这里指定的是高度维度可变，如果输入时通道数发生变化就会导出失败

一般来讲pytorch模型转onnx，主要包括三部分，在torch中实现该算子，有对应的符号函数连接torch算子和onnx算子，在onnx中该算子有实现。

查看torch和onnx中算子的对应情况，在 PyTorch 中，和 ONNX 有关的定义全部放在 torch.onnx目录中，如果找不到这个目录可以参考这里
可以在对应的算子集版本中查找你想知道的算子名称即可，在symbolic_fn中查看其与onnx算子的映射情况。onnx算子文档在这里

有些时候，我们希望模型在导出至 ONNX 时有一些不同的行为模型在直接用 PyTorch 推理时有一套逻辑，而在导出的ONNX模型中有另一套逻辑。比如，我们可以把一些后处理的逻辑放在模型里，以简化除运行模型之外的其他代码。torch.onnx.is_in_onnx_export()可以实现这一任务，该函数仅在执行 torch.onnx.export()时为真

一般来说推荐模型的输入为张量而不是python常量，如果在跟踪的过程中涉及到张量与常量的转换可能会导致导出失败，另一方面，我们也可以利用这个性质，在保证正确性的前提下令模型的中间结果变成常量。这个技巧常常用于模型的静态化上，即令模型中所有的张量形状都变成常量。

算子支持

PyTorch的的代码主要由C10、ATen(PyTorch 的 C++ 张量运算库)、torch三大部分组成的
这一节主要讲了三点

1. 如何用register_op为ATen算子添加符号函数(PyTorch 算子类的一个静态方法)
2. register_custom_op_symbolic 为 TorchScript 算子补充注册符号函数
3. 添加C++拓展然后用torch.autograd.function封装，编写符号函数进行映射，然后就可以在torch中调用

1.使用情况:算子在 ATen 中已经实现了, ONNX 中也有相关算子的定义，但是相关算子映射成 ONNX 的规则没有写。在这种情况下，我们只需要为 ATen 算子补充描述映射规则的符号函数就行了。

import torch
import numpy as np
import onnxruntime
from torch.onnx.symbolic_registry import register_op

class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        return torch.asinh(x)
#绑定onnx算子
def asinh_symbollic(g, input, *, out = None): #asinh的符号函数,从除g以外的第二个输入参数开始，其输入参数应该严格对应pyi文件中的函数接口：
    return g.op('Asinh', input) #完成了 ONNX 算子的定义。其中，第一个参数"Asinh"是算子在 ONNX 中的名称。

#把符号函数 asinh_symbolic 绑定到ATen算子 asinh 上
register_op('asinh', asinh_symbollic, '', 9)
#第一个参数是目标 ATen 算子名，第二个是要注册的符号函数。第三个参数是算子的“域”，对于普通 ONNX 算子，直接填空字符串即可。第四个参数表示向哪个算子集版本注册。

model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'asinh.onnx')


#test op 
torch_output = model(input).detach().numpy()

sess = onnxruntime.InferenceSession('asinh.onnx')
ort_output = sess.run(None, {'0':input.numpy()})[0]

assert np.allclose(torch_output, ort_output)

添加C++拓展

// my_add.cpp

#include 

//算子的实现

torch::Tensor my_add(torch::Tensor a, torch::Tensor b) 
{
    return 2*a + b;
}

//PYBIND11_MODULE 来为 C++ 函数提供 Python 调用接口
PYBIND11_MODULE(my_lib, m) 
{
    m.def("my_add", my_add);
} //my_lib 是我们未来要在 Python 里导入的模块名。双引号中的 my_add 是 Python 调用接口的名称，这里我们对齐 C++ 函数的名称，依然用 "my_add"这个名字。

//之后，我们编写 Python 代码并命名为 "setup.py"，来编译刚刚的 C++ 文件

setup.py用来编译上面的C++文件

from setuptools import setup
from torch.utils import cpp_extension
#这段代码使用了 Python 的 setuptools 编译功能和 PyTorch 的 C++ 拓展工具函数，可以编译包含了 torch 库的 C++ 源文件。

setup(name = 'my_add', ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])], 
        cmdclass= {'build_ext':cpp_extension.BuildExtension})
#这里我们需要填写的只有模块名和模块中的源文件名。我们刚刚把模块命名为 my_lib，而源文件只有一个 my_add.cpp

之后切换到setup文件夹下，执行

python setup.py develop 

即可编译

对该算子进行封装并调用

import torch
import my_lib

class MyAddFunction(torch.autograd.Function): #用 torch.autograd.Function 来封装算子的底层调用,把算子封装成 Function
    @staticmethod
    def forward(ctx, a, b):
        return my_lib.my_add(a,b) 

    @staticmethod
    def symbolic(g, a, b): #用 g.op() 定义了三个算子：常量、乘法、加法
        two = g.op('Constant', value_t  = torch.tensor([2])) #把 PyTorch 张量的值传入 value_t 参数中
        a = g.op('Mul', a, two)
        return g.op('Add', a, b)

my_add = MyAddFunction.apply #apply是torch.autograd.Function 的一个方法，这个方法完成了 Function 在前向推理或者反向传播时的调度
#这里 my_add 的地位，和 PyTorch 的 asinh, interpolate, conv2d等原生函数是类似的

class MyAdd(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return my_add(a, b)

model = MyAdd()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, (input,input), 'my_add.onnx')
torch_output = model(input, input).detach().numpy()

import numpy as np
import onnxruntime

sess = onnxruntime.InferenceSession('my_add.onnx')
ort_output = sess.run(None, {'a':input.numpy(), 'b':input.numpy()})[0]

assert np.allcloese(torch_output, ort_output)

转换为onnx就先到这里，之后进行onnx到APK这一步骤