三、YOLOv5模型转onnx

前面说完YOLOv5的训练，也进行了相应的测试，接下来就是对训练好的pt模型转为onnx模型！
在YOLOv5的git项目里有自带的一个onnx_export.py文件，运行该文件，即可将pt模型转为onnx模型，但是里面也有很多坑！！！！
先来看onnx_export.py的代码：

"""Exports a pytorch *.pt model to *.onnx format

Usage:
    import torch
    $ export PYTHONPATH="$PWD" && python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
"""

import argparse

import onnx

from models.common import *

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='./weights/last.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    opt = parser.parse_args()
    print(opt)

    # Parameters
    f = opt.weights.replace('.pt', '.onnx')  # onnx filename
    img = torch.zeros((opt.batch_size, 3, *opt.img_size))  # image size, (1, 3, 320, 192) iDetection

    # Load pytorch model
    # google_utils.attempt_download(opt.weights)
    model = torch.load(opt.weights,  map_location=torch.device('cpu'))['model']
    model.eval()
    model.fuse()

    # Export to onnx
    model.model[-1].export = True  # set Detect() layer export=True Ture会使detect层不加入onnx
    _ = model(img)  # dry run
    torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['images'],
                      output_names=['output'])  # output_names=['classes', 'boxes']

    # Check onnx model
    model = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model)  # check onnx model
    print(onnx.helper.printable_graph(model.graph))  # print a human readable representation of the graph
    print('Export complete. ONNX model saved to %s\nView with https://github.com/lutzroeder/netron' % f)

需要修改onnx_export.py代码中的几个地方：
pt模型默认输入size为640*640，这里转onnx时将输入减半，是由于在调用ncnn时需要重写后处理部分代码，会有多个for循环，为了提速，将输入大小减半，具体的输入应该是可以自己设置的，只要是32的倍数即可，但是本文没有尝试，感兴趣的可自行尝试！
注意！YOLOv5在转onnx的时候，有个巨坑，就是在网络模型中有个Focus机制，其中有一个tensor的切片操作，这个操作在转ncnn的时候不被支持，需要对这个模型的切片操作进行修改！
具体的修改是在models文件夹下的common.py文件，其原始代码如下：

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1):
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, 1)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        # return self.conv(torch.cat([x, x, x, x], 1))  
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

我们要将上述代码做如下修改：
将切片操作修改之后，转ncnn的时候才能成功，当然，这么修改模型对精度是有一定影响的！如果你足够牛，可以去ncnn里自定义实现切片操作！！！！

到这里为止，就已经将pt模型转为了onnx模型，可以说已经脱离了pytorch了，但是转ncnn还需要对模型做一个Simplifier操作，因为转出的onnx模型还有许多冗余，这在ncnn里也是不支持的，避免转ncnn时各种报错，就先操作一下把！

直接在pycharm里打开Terminal，在里面输入以下代码：

python3 -m onnxsim input_onnx_model out_onnx_model

然后回车即可，其中input_onnx_model是你需要简化的onnx的路径，out_onnx_model就是输出简化后模型的路径！

转onnx到这一步就算结束了，接下来就是onnx转ncnn了！

四、onnx转ncnn

如果是按照上面的步骤将yolov5转为onnx了，那么接下来转ncnn应该是非常顺利的，本文亲测有效！如若不是，过程中多半会出问题，就自求多福吧！
ncnn编译的问题可以去看我的另一篇博客，里面有详细介绍，只要ncnn成功编译了，通过终端cd到ncnn/build/tools/onnx/路径下，里面有一个编译好的onnx2ncnn文件，如下：
这时只需要在终端中输入：

./onnx2ncnn model.onnx model.param model.bin

回车即可，其中model.onnx即是需要转为ncnn的onnx模型，model.param和model.bin即为转为ncnn后输出的两个文件，注意这两个文件的顺序千万不能写反了！！！

如果什么都没输出，那恭喜成功了！如果遇到下面这种情况：
多半是没做Simplifier操作，若果做了还是不行，那自求多福吧！！！

最终生成的ncnn文件如下：
恭喜你，成功转为ncnn模型了！

更多技术欢迎加入交流：320297153