自己动手写深度学习框架（pytorch转ncnn）

pytorch本身是我们进行ai训练的一个框架，但是训练得到的model最终是要部署到嵌入式板子上面的。这中间就有两种办法，一种是pytorch直接生成onnx文件，然后用ncnn自带的onnx2ncnn将onnx转成我们想要的文件。另外一种方法，就是先生成pt文件，接着用pnnx工具把pt生成我们想要的ncnn文件。

1、pytorch-》onnx-》ncnn

这种方法比较简单，只是在onnx转ncnn这一步的时候常常会有告警提示。首先，我们先安装onnxscript库，用pytorch生成onnx文件，

复制代码

pip3.10.exe install onnxscript -i https://pypi.doubanio.com/simple

接着继续生成onnx文件，

复制代码

dummy_input = torch.randn(1, n_in)
torch.onnx.export(
    model,
    dummy_input,
    "my_model.onnx",
    export_params=True,
    opset_version=18,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

有了onnx文件，继续借助于ncnn开发包里面的onnx2ncnn可执行文件，就可以将onnx换成我们想要的的param文件了。命令行是这样的，

复制代码

onnx2ncnn my_module.onnx my_module.param my_module.bin

只不过，在转换的时候，大部分时候都会有告警，提醒大家转换的过程中可能有数据丢失，建议大家用pnnx转换。

2、pt-》pnnx转换

和之前的方法比较，这种方法需要先安装pnnx，用pip310安装即可。安装完了，pnnx可执行文件就安装在scripts目录下，

复制代码

pip3.10.exe install pnnx  -i https://pypi.doubanio.com/simple

接下来，就是准备pt文件，尽量按照这个模板来，

复制代码

import torch

model.eval() # generate pt, then use pnnx & pt to generate ncnn files
dummy_input = torch.randn(1, n_in)
mod = torch.jit.trace(model, x)
mod.save("my_model.pt")

执行这个模板之后，就可以用pnnx生成对应的param文件和bin文件了。

复制代码

C:\Python310\Scripts\pnnx.exe my_model.pt

3、ncnn需要的两个文件

ncnn要运行起来，最主要的，就是param文件和bin文件。其中param文件描述model，bin则是model里面的参数内容。这里看一下param文件里面是什么样的，

复制代码

7767517
5 5
Input                    in0                      0 1 in0
Gemm                     gemm_0                   1 1 in0 1 10=4 2=0 3=1 4=0 5=1 6=1 7=10 8=5 9=10
ReLU                     relu_2                   1 1 1 2
Gemm                     gemm_1                   1 1 2 3 10=4 2=0 3=1 4=0 5=1 6=1 7=10 8=1 9=5
Sigmoid                  sigmoid_3                1 1 3 out0

bin文件就是纯二进制，这里就不多做展示了。