當前位置:網站首頁>將pytorch訓練的模型轉為ncnn模型

將pytorch訓練的模型轉為ncnn模型

2022-01-27 23:35:37 樊城

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:
    在這裏插入圖片描述
    在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

版權聲明
本文為[樊城]所創,轉載請帶上原文鏈接,感謝
https://cht.chowdera.com/2022/01/202201272335367678.html

隨機推薦