【pytorch 转 onnx】pytorch-onnx
【onnx 转 ncnn】onnx-ncnn
【ncnn 加载模型】ncnn-load
一、python安装依赖项
pip install onnxruntime onnx opencv-python
二、创建模型并训练,加载模型参数并输出onnx
#### pytorch 转 onnx
import torch
import torch.onnx
from model.efficientnet_model import EfficientNet_FV
model = EfficientNet_FV(num_classes=4800).cuda()
path_checkpoint = "..path\\best.ckpt"
checkpoint = torch.load(path_checkpoint)
model.load_state_dict(checkpoint) # 加载训练好的模型
model.eval()
##### 生成 onnx 文件
x = torch.randn(1, 3, 128, 128).cuda()
with torch.no_grad():
torch.onnx.export(
model, x,
"EfficientNet_b0.onnx",
opset_version=11,
export_params=True,
input_names=["input"],
output_names=["output1"],
)
其中,torch.onnx.export 是 PyTorch 自带的把模型转换成 ONNX 格式的函数。前三个参数分别是要转换的模型、模型的任意一组输入、导出的 ONNX 文件的文件名。转换模型时,需要原模型和输出文件名是很容易理解的,但为什么需要为模型提供一组输入呢?这就涉及到 ONNX 转换的原理了。从 PyTorch 的模型到 ONNX 的模型,本质上是一种语言上的翻译。直觉上的想法是像编译器一样彻底解析原模型的代码,记录所有控制流。但前面也讲到,我们通常只用 ONNX 记录不考虑控制流的静态图。因此,PyTorch 提供了一种叫做追踪(trace)的模型转换方法:给定一组输入,再实际执行一遍模型,即把这组输入对应的计算图记录下来,保存为 ONNX 格式。export 函数用的就是追踪导出方法,需要给任意一组输入,让模型跑起来。我的测试图片是三通道,3*128*128大小的,这里也构造一个同样形状的随机张量。
如果上述代码运行成功,目录下会新增一个"EfficientNet_b0.onnx"的 ONNX 模型文件。我们可以用下面的脚本来验证一下模型文件是否正确。
import onnx
onnx_model = onnx.load("EfficientNet_b0.onnx")
try:
onnx.checker.check_model(onnx_model)
except Exception:
print("Model incorrect")
else:
print("Model correct")
可以使用 Netron (开源的模型可视化工具)来可视化 ONNX 模型。
三、onnx2ncnn
关于onnx2ncnn:这个是重点,得到了onnx模型,就告别了pytorch的框架了,现在需要将onnx格式的模型转成ncnn能够加载的.bin和.param。用前面编译好的ncnn带的工具,如果ncnn编译成功的话:ncnn/build/tools/onnx/下面有一个编译好的onnx2ncnn。只需要输入以下命令,(model是onnx名字),如果什么都没输出,那就成功了。一般不建议将除网络以外的部分打包到onnx中。
虽然成功生成了.param和.bin,但是用ncnn 来load模型的时候一般都会失败!因为目前ncnn对于某一些操作还不支持,所以这里还有引入一个工具:ONNX Simplifier。
GitHub - daquexian/onnx-simplifier: Simplify your onnx model
pip install onnxsim
onnxsim input_onnx_model output_onnx_model
四、ncnn加载模型
参考:https://blog.csdn.net/Enchanted_ZhouH/article/details/106063552