Pytorch模型部署：torch-onnx-ncnn

VIP文章亿丶Bug

已于 2022-08-29 21:12:14 修改

阅读量1.9k

点赞数 1

分类专栏：部署文章标签： pytorch 人工智能 python

于 2022-08-22 15:47:17 首次发布

本文链接：https://blog.csdn.net/my_warning/article/details/125187414

版权

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

学习内容：

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的生成主要包括三步：

使用 trace 的方式先生成一个 TorchScipt 模型，如果你转换的本身就是 TorchScript 模型，则可以跳过这一步。
使用许多 pass 对 1 中生成的模型进行变换，其中对 ONNX 导出最重要的一个 pass 就是ToONNX，这个 pass会进行一个映射，将 TorchScript 中prim、aten空间下的算子映射到onnx空间下的算子。
使用 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('.'