将pytorch训练的模型转为ncnn模型

1、训练自己的pytorch模型

训练完成后,生成.pth后缀模型,此时要转为可部署的ncnn模型需要经历一些过程,如下:.pth -> .onnx -> .pram 和 .bin模型

2、.pth转.onnx

2.1 避坑

有时候会出现onnx库不能将某些模块转换出来,如mobilenet中的hardswish激活函数就不能被转换,原因是由于pytorch的版本低,导致无法转换,可以升级pytorch的版本。

2.2 代码

以下代码为将.pth文件转为.onnx文件

import torch
import torchvision
import torch.nn as nn
import collections
from model_v3 import mobilenet_v3_large
# from torchsummary import summary


model_weight_path = "./MobileNetV3_large230.pth"   #自己的pth文件路径
out_onnx = './MobileNetV3_large230.onnx'           #保存生成的onnx文件路径
model = mobilenet_v3_large( num_classes=2)   # 加载自己的的网络
model.load_state_dict(torch.load(model_weight_path)) #加载自己的pth文件
model.eval()

x = torch.randn(1, 3, 224, 224)
#define input and output nodes, can be customized
input_names = ["input"]
output_names = ["output"]
#convert pytorch to onnx
torch_out = torch.onnx.export(model, x, out_onnx, input_names=input_names, output_names=output_names)

3、.onnx转ncnn模型

这一步过程相对来说复杂一点。

3.1 下载cmake3.21.0\protobuf3.4.0\ncnn
  • cmake需要添加环境变量(将解压后cmake-3.21.0文件夹点开进入bin将该路径添加至环境变量)
  • vs2019编译protobuf:
    将pytorch训练的模型转为ncnn模型_第1张图片
    在VS2019的X64命令行下执行以下命令:(首先cd到protobuf-3.4.0的解压后的root目录)
    依次输入以下指令:
> cd 
> mkdir build-vs2019
> cd build-vs2019
> cmake -G"NMake Makefiles" -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=%cd%/install -Dprotobuf_BUILD_TESTS=OFF -Dprotobuf_MSVC_STATIC_RUNTIME=OFF ../cmake
> nmake
> nmake install

结束后在protobuf文件夹下生成build-vs2019文件夹及相应文件

  • ncnn编译:
    同样在在VS2019的X64命令行下执行以下命令:(首先cd到ncnn的解压后的root目录)
> cd <ncnn-root-dir>
> mkdir -p build-vs2019
> cd build-vs2019
> cmake -G"NMake Makefiles" -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=%cd%/install -DProtobuf_INCLUDE_DIR=<protobuf-root-dir>/build-vs2019/install/include -DProtobuf_LIBRARIES=<protobuf-root-dir>/build-vs2019/install/lib/libprotobuf.lib -DProtobuf_PROTOC_EXECUTABLE=<protobuf-root-dir>/build-vs2019/install/bin/protoc.exe ..
> nmake
> nmake install

注意:" protobuf-root-dir "为protobuf根目录,<>应该删除。
编译成功后,可以在build-vs2019\tools\onnx文件夹下找到onnx2ncnn可执行文件。

参考博客:
https://blog.csdn.net/baidu_33512336/article/details/114301368

你可能感兴趣的:(ncnn,python,深度学习,pytorch)