pytorch2ncnn

借助 onnx 实现 pytorch 到 ncnn 的转换

1. 安装 onnx

pip3 install onnx

2. pytorch model > onnx

from models.pfld import PFLDInference, AuxiliaryNet
from torch.autograd import Variable
import onnx
import torch

checkpoint = torch.load("./checkpoint/snapshot/checkpoint.pth.tar", map_location='cpu')
pfld_backbone = PFLDInference()
pfld_backbone.load_state_dict(checkpoint['pfld_backbone'])

dummy_input = Variable(torch.randn(1, 3, 112, 112)) 
torch.onnx.export(plfd_backbone, dummy_input, "output/plfd.onnx")

model = onnx.load('output/pfld.onnx')

3. 安装ncnn

(1) 安装 protobuf

unzip protobuf-all-3.7.1.zip
cd protobuf-all-3.7.1
./configure --prefix=/usr/local
make -j4
make check -j4
sudo make install
sudo ldconfig

(2) 安装 opencv, opencv 这里用的是2.4.13.6

## 源码编译
unzip opencv-2.4.13.6.zip
cd opencv-2.4.13.6
mkdir build
cd build
cmake -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/usr/local ..
make -j8
sudo make install
sudo ldconfig

(3) 安装依赖项

sudo apt-get install build-essential
sudo apt-get install cmake git libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev
sudo apt-get install  libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libjasper-dev libdc1394-22-dev

(4) 把ncnn的源码clone下来,我们这里主要想使用tools下面的onnx转换工具。

这里作一点修改,pytorch1.0之后支持0维的张量,这在ncnn转换中会出现问题,修改onnx2ncnn.cpp中Constant和MemoryData的转换,有2处

if (M.dims_size() == 1) {
   fprintf(pp, " 0=%d", (int)M.dims(0));
} else if (M.dims_size() == 2) {
    fprintf(pp, " 0=%d", (int)M.dims(1));
    fprintf(pp, " 1=%d", (int)M.dims(0));
} else if (M.dims_size() == 3) {
    fprintf(pp, " 0=%d", (int)M.dims(2));
    fprintf(pp, " 1=%d", (int)M.dims(1));
    fprintf(pp, " 2=%d", (int)M.dims(0));
} else if (get_tensor_proto_data_size(M)==1) {
    // scalar tensor!!! 
    fprintf(pp, " 0=1");
}

(3) 编译 ncnn

cd ncnn
mkdir -p build
cd build
cmake ..
make -j4

这样在 build/tools/onnx/ 目录下就有转换工具 onnx2ncnn 了

4. 简化一些onnx 模型

先别急着转换,onnx转换模型时有一些冗余,我们用工具简化一些onnx模型

pip3 install onnx-simplifier
python3 -m onnxsim pfld.onnx pfld-sim.onnx

5. onnx > ncnn

cd ~/xxx/ncnn/build/tools/onnx/
./onnx2ncnn pfld-sim.onnx pfld-sim.param pfld-sim.bin

现在就生成了两个文件可以供使用: pfld-sim.binpfld-sim.param

6. 编写一个可以供调用的ncnn文件进行调用即可。

错误信息:

Q: Python中Import Error: no module named ‘past’错误以及解决方法

pip3 install future

本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!