背景:
- 因为所里面大多数同学在研究论文时用的比较多的还是pytorch,所以考虑在pc端的部署采用libtorch ,当然
libtorch比较新,所以还是要谨慎地采用,因为看到官方只用resnet等简单的网络做例子时就要有被坑的心里准备,但是,,项目紧急,不管了。。 - 只要懂得其中一些规则,尤其是jit 的规则还是挺好弄的,但是本人也是第一次真正接触到jit,所以也踩了很多坑
- 如果熟悉了里面的python和torch script的转换规则,那么就会轻松很多
- 本文是以SSD(single shot multibox object detector)的python版本为原型,转换成torch script c++版本调用,这是原文的链接https://github.com/amdegroot/ssd.pytorch
note !:
- 本文采用的是Annotation方式,不是trace
- 因为上面gayhub的ssd 里面有些python操作是torch script 不支持的,所以我把最后一步的操作分解了,同时原文的nms操作好像是直接python操作的,没有cuda并行版本(可能本人没找到),所以本文结合该ssd修改GitHub - gdlg/pytorch_nms: CUDA implementation of NMS for PyTorch的把nms变成多图多类的nms并行化操作
- 关于速度问题,如果只是单张图片且只有一类,那么推荐nms在cpu上操作,那样nms会比在gpu上nms操作快,可能论单核的话还是cpu强一些
本文涉及到的一些知识:
- 一些ssd的知识,但是没有上篇讲mask rcnn的仔细。因为ssd稍微简单一些,可能是比mask rcnn 更‘端到端化一些’
- libtorch c++一些常规操作
- nms的并行算法原理(这个可能有点难理解)
- cuda的一些常规操作(因为本人也是刚接触cuda,所以可能修改后的nms并不是常规操作)
环境配置:
- 系统: windows 10
- libtorch 版本: libtorch-win-shared-with-deps-1.2.0.zip
- IDE : QT creator 5.8.0 编译器:Qt 5.8 64-bit for Desktop (MSVC 2015)
- GPU: nvidia 1080ti(11G)
1 pytorch模型转libtorch(script)模型
先下载pytorch ssd代码GitHub - amdegroot/ssd.pytorch: A PyTorch Implementation of Single Shot MultiBox Detector
因为我们是对训练好的pytorch模型进行转换,所以我们可以只要测试的代码就行了,所以我们可以直接对SSD文件夹主目录下的test.py进行修改就行了,修改前先进行备份(以后涉及到修改的文件都可以备份),假设命名为convert2libtorch.py,之后我们对该文件进行修改,
修改开始(按照官网的教程来Loading a TorchScript Model in C++ — PyTorch Tutorials 1.10.1+cu102 documentation,主要是annotation方式,实际上很多复杂点的网络好多control flow):
1. 我们主要对convert2libtorch.py 的test_voc函数进行修改,
- 先添加torch script用到的包,同时可以添加cv2包(因为测试时可能会用到)
- 其他包比如
- 因为我这里没有下载voc,所以里面很多路径是没有的,但是这样运行会出错,所以我这边的文件改成如下
- 如果你是下载过数据集并且完整跑过这些代码,那么上面的可以不改
下载方法:
源码 baidu链接 xxx +1KMFWV95HQ3azBZUuqM_VDA tq码:rfu1 ,xxx是百度盘的前缀,防吞
下载文件放资源里了
libtorch主要问题放在这里:libtorch 问题汇总_captain_CasonCai的博客_libtorch 内存 泄露