3、YOLOv5模型转onnx
前面说完YOLOv5的训练,也进行了相应的测试,接下来就是对训练好的pt模型转为onnx模型!
在YOLOv5的git项目里有自带的一个onnx_export.py文件,运行该文件,便可将pt模型转为onnx模型,可是里面也有不少坑!!!!
先来看onnx_export.py的代码:python
"""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文件,其原始代码以下:git
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里自定义实现切片操做!!!!github
到这里为止,就已经将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就是输出简化后模型的路径!ui
转onnx到这一步就算结束了,接下来就是onnx转ncnn了!google
4、onnx转ncnn
若是是按照上面的步骤将yolov5转为onnx了,那么接下来转ncnn应该是很是顺利的,本文亲测有效!如若不是,过程当中多半会出问题,就自求多福吧!
ncnn编译的问题能够去看个人另外一篇博客,里面有详细介绍,只要ncnn成功编译了,经过终端cd到ncnn/build/tools/onnx/路径下,里面有一个编译好的onnx2ncnn文件,以下:
这时只须要在终端中输入:spa
./onnx2ncnn model.onnx model.param model.bin
回车便可,其中model.onnx便是须要转为ncnn的onnx模型,model.param和model.bin即为转为ncnn后输出的两个文件,注意这两个文件的顺序千万不能写反了!!!code
若是什么都没输出,那恭喜成功了!若是遇到下面这种状况:
多半是没作Simplifier操做,若果作了仍是不行,那自求多福吧!!!orm
最终生成的ncnn文件以下:
恭喜你,成功转为ncnn模型了!