Windows系统下把PyTorch模型转为ncnn模型流程

文章目录

  • 前言
  • 一、大概流程及ONNX模型简介
  • 二、环境配置
    • 1、需要安装的软件
    • 2、在Windows系统下构建ncnn环境
    • 3、VS2015配置
  • 三、步骤
    • 1、把PyTorch模型(.pth文件)转为onnx模型(.onnx文件)
    • 2、简化onnx模型
    • 3、生成ncnn模型
    • 4、使用VS编译ncnn模型
  • 总结


前言

  最近需要部署深度学习模型,选用了腾讯的ncnn框架,也就是要把训练好的PyTorch模型(.pth文件)转换为ncnn模型(.param和.bin文件),这次是在Windows 10系统部署,中间遇到了很多坑,还好参考各方大佬的教程才得以成功,在此记录一下,以便能帮助后来的学习者。

一、大概流程及ONNX模型简介

  首先要把PyTorch模型(.pth文件)转为onnx模型(.onnx文件),然后简化onnx模型,接着才能把onnx模型转为ncnn模型(.param和.bin文件)。
  onnx(开放式神经网络交换)是一个中间表达格式框架,onnx是一种针对机器学习所设计的开放式的文件格式,用于存储训练好的模型。它使得不同的人工智能框架(如PyTorch, MXNet)可以采用相同格式存储模型数据并交互, onnx的规范及代码主要由微软,亚马逊 ,Facebook 和IBM等公司共同开发,以开放源代码的方式托管在GitHub上。

二、环境配置

  配置环境是模型转换的前提,遇到的坑基本在配置环境上了,详细描述如下:

1、需要安装的软件

  需要安装的软件为:Visual Studio(我用的2015版本)、OpenCV、Cmake。
  为什么需要使用这三个软件?VS是美国微软公司的开发工具包系列产品,在该项目中有提供GCC、编译运行ncnn模型的C++程序的作用;CMake 是一个跨平台的,开源的构建系统,CMake可以通过CMakeLists.txt文件来产生特定平台的标准的构建文件,例如:为Unix平台生成makefiles文件(使用GCC编译),为Windows MSVC生成 projects/workspaces(使用VS IDE编译)或Makefile文件(使用nmake编译);OpenCV是一个基于BSD许可(开源)发行的跨平台计算机视觉和机器学习软件库,学习计算机视觉基本上离不开OpenCV。
  Visual Studio 2015下载链接如下:MSDN,我告诉你,Visual Studio分为三个版本:Community、Enterprise、Professional, 使用Community足够了,把下载相应版本的链接复制粘贴到迅雷上下载;
  OpenCV下载链接如下:OpenCV下载链接,我下载的4.5.2版本,下载.exe,安装和配置可见链接OpenCV安装配置链接;
  cmake下载链接如下:cmake下载链接,我下载的3.20.1版本,如何安装可参考链接:Windows下载安装cmake。

2、在Windows系统下构建ncnn环境

  可参考官方链接:官方链接
  构建步骤如下:
  (1)下载ncnn,Git Bash下指令如下:

$ git clone https://github.com/Tencent/ncnn.git
$ cd ncnn
$ git submodule update --init                      #可以下载工程里的一些子工程

  (2)重点:接下来的步骤使用VS2015的本地工具命令提示符,而不是cmd和Git Bash,在左下角开始菜单中,找到VS的本地工具命令提示符,如下图所示:
在这里插入图片描述
  (3)下载protobuf,下载链接如下:下载链接,解压到指定文件夹,最好和ncnn放在同一个目录。
  为什么要下载protobuf文件?protobuf用于转换模型,protobuf(Google Protocol Buffers)是Google提供一个具有高效的协议数据交换格式工具库(类似Json),Protobuf 提供了C++、java、python语言的支持,提供了windows(proto.exe)和linux平台动态编译生成proto文件对应的源文件。
  (4)在VS2015的本地工具命令提示符下构建protobuf(相当于用VS2015编译protobuf),指令如下:

> cd <protobuf-root-dir>                   #是指protobuf文件夹的根目录
> mkdir build_vs2015
> cd build_vs2015
> cmake -G"NMake Makefiles" -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=%cd%/install -Dprotobuf_BUILD_TESTS=OFF -Dprotobuf_MSVC_STATIC_RUNTIME=OFF ../cmake
> nmake                                   #编译cmake生成的Makefile文件
> nmake install                           #安装操作,把生成的文件复制到对应的目录中,并修改环境变量等。

  cmake指令解释如下:cmake -G的G是Generator的意思,该语句可以设置IDE,比如VS;CMAKE_BUILD_TYPE是要编译的类型,一般的选择有debug,release;CMAKE_INSTALL_PREFIX是安装路径,就是执行make install时的安装目录(windows系统中对应nmake install),后面以Dprotobuf开头的是开关;最后一个表示cmake用到的CMakeLists.txt文件所在目录。
  关于make/makefile/cmake/nmake的解释,可见链接:解释。

  cmake成功后的界面如下:
Windows系统下把PyTorch模型转为ncnn模型流程_第1张图片
  nmake成功后的界面如下:
Windows系统下把PyTorch模型转为ncnn模型流程_第2张图片
  nmake install成功后的界面如下:
Windows系统下把PyTorch模型转为ncnn模型流程_第3张图片

  (5)下载Vulkan SDK(可选),下载链接如下:Vulkan SDK下载链接,使用Valkan就可以用GPU进行推理,我使用CPU推理,所以没有下载Vulkan SDK。
  (6)在VS2015的本地工具命令提示符下构建ncnn,指令如下:

> cd <ncnn-root-dir>                   #是指ncnn的根目录
> mkdir build
> cd build
> cmake -G"NMake Makefiles" -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=%cd%/install -DProtobuf_INCLUDE_DIR=<protobuf-root-dir>/build/install/include -DProtobuf_LIBRARIES=<protobuf-root-dir>/build/install/lib/libprotobuf.lib -DProtobuf_PROTOC_EXECUTABLE=<protobuf-root-dir>/build/install/bin/protoc.exe -DNCNN_VULKAN=OFF ..
> nmake
> nmake install

  把cmake命令下DProtobuf开头命令的路径改成自己protobuf所在路径。

3、VS2015配置

  (1)新建VS2015工程,在上面一栏,打开View->Other Windows->Property Manager,在侧面会出现:
Windows系统下把PyTorch模型转为ncnn模型流程_第4张图片
  双击Release、x64版本的Microsoft.Cpp.x64.user,出现如下配置界面:
Windows系统下把PyTorch模型转为ncnn模型流程_第5张图片
  接下来主要是添加一些目录,首先添加VC++的include目录,如下图:
Windows系统下把PyTorch模型转为ncnn模型流程_第6张图片
  然后添加VC++的lib目录,如下图:
Windows系统下把PyTorch模型转为ncnn模型流程_第7张图片
  然后添加VC++的lib WinRT目录(Windows运行库目录),如下图:
Windows系统下把PyTorch模型转为ncnn模型流程_第8张图片
  然后在linker的input中添加Additional Dependencies(附加依赖项),lib具体名称可查阅自己文件下的相应名称,如下图:
Windows系统下把PyTorch模型转为ncnn模型流程_第9张图片
  至此已将环境全部配置完毕。

三、步骤

1、把PyTorch模型(.pth文件)转为onnx模型(.onnx文件)

  我们需要训练好PyTorch模型,生成.pth文件,才能进行这一步,其实转换十分简单,仅需几行代码,我截取了自己转换时的代码,大家对应自己项目修改即可,代码如下:

self.net = RetinaFace(cfg=self.cfg, mode='eval').eval()                             #实例模型
state_dict = torch.load(self.model_path,map_location=lambda storage, loc: storage)  #将生成的.pth文件加载到内存中,我使用CPU推理GPU训练的模型,所以需要加上后两个参数
self.net.load_state_dict(state_dict)                                                #将加载到内存中的.pth文件加载到模型中
x = torch.rand(1, 3, 224, 224)
torch_out = torch.onnx.export(self.net, x, "retinaface.onnx")                      #将PyTorch模型转换为onnx模型
print('Finished!')

  转换函数参数说明如下:

torch.onnx.export(network, x, "retinaface.onnx", export_params=True, verbose=True,input_names=input_names, output_names=output_names)
network :模型
x:模型输入
"retinaface.onnx" :模型保存名
export_params:是否保存参数
verbose:查看模型信息
input_name:定义输入层名
output_name:定义输出层名

2、简化onnx模型

  如果直接将转换后的onnx模型转为ncnn模型,会出现如下错误:
Windows系统下把PyTorch模型转为ncnn模型流程_第10张图片
  因此需要简化onnx模型,首先安装简化工具,指令如下:

pip install onnx-simplifier

  然后简化onnx模型,指令如下:

python -m onnxsim retinaface.onnx retinaface-sim.onnx

  即简化成功,简化后的onnx模型为retinaface-sim.onnx。

3、生成ncnn模型

  将retinaface-sim.onnx模型复制粘贴到ncnn-root-dir\build\tools\onnx文件夹下面,使用如下指令即可转换成ncnn模型:

onnx2ncnn retinaface-sim.onnx retinaface.param retinaface.bin

  ncnn模型有两个文件(.param和.bin),.param可以理解为网络的配置文件,.bin可以理解为网络的参数(各种权重)文件。

4、使用VS编译ncnn模型

  在上一步配置好的VS工程下新建一个cpp文件,根据ncnn模型架构写cpp代码,然后编译运行(快捷键为Ctrl+F5),注意运行Release X64版本,如下图:
在这里插入图片描述

  如下图所示,进入生成.exe文件的目录,输入对应的指令(参数为传入main函数的参数),即可成功运行:
Windows系统下把PyTorch模型转为ncnn模型流程_第11张图片

总结

  1、配置环境的步骤对于新手来说相当繁琐,但是在这个过程中也收获了很多;
  2、需要根据模型架构写对应的cpp代码,这个也挺有难度。

你可能感兴趣的