ssd pytorch转 libtorch c++实现 (nms采用cuda并行实现)

背景:

  1. 因为所里面大多数同学在研究论文时用的比较多的还是pytorch,所以考虑在pc端的部署采用libtorch ,当然
    libtorch比较新,所以还是要谨慎地采用,因为看到官方只用resnet等简单的网络做例子时就要有被坑的心里准备,但是,,项目紧急,不管了。。
  2. 只要懂得其中一些规则,尤其是jit 的规则还是挺好弄的,但是本人也是第一次真正接触到jit,所以也踩了很多坑
  3. 如果熟悉了里面的python和torch script的转换规则,那么就会轻松很多
  4. 本文是以SSD(single shot multibox object detector)的python版本为原型,转换成torch script c++版本调用,这是原文的链接https://github.com/amdegroot/ssd.pytorch

note !:

  1. 本文采用的是Annotation方式,不是trace

  2. 因为上面gayhub的ssd 里面有些python操作是torch script 不支持的,所以我把最后一步的操作分解了,同时原文的nms操作好像是直接python操作的,没有cuda并行版本(可能本人没找到),所以本文结合该ssd修改https://github.com/gdlg/pytorch_nms的把nms变成多图多类的nms并行化操作

  3. 关于速度问题,如果只是单张图片且只有一类,那么推荐nms在cpu上操作,那样nms会比在gpu上nms操作快,可能论单核的话还是gpu强一些
     

本文涉及到的一些知识:

  1. 一些ssd的知识,但是没有上篇讲mask rcnn的仔细。因为ssd稍微简单一些,可能是比mask rcnn 更‘端到端化一些’
  2. libtorch c++一些常规操作
  3. nms的并行算法原理(这个可能有点难理解)
  4. cuda的一些常规操作(因为本人也是刚接触cuda,所以可能修改后的nms并不是常规操作)

环境配置:

  1. 系统: windows 10
  2. libtorch 版本: libtorch-win-shared-with-deps-1.2.0.zip
  3. IDE : QT creator  5.8.0  编译器:Qt 5.8 64-bit for Desktop (MSVC 2015) 
  4. GPU: nvidia 1080ti(11G)

1 pytorch模型转libtorch(script)模型

  • 先下载pytorch ssd代码https://github.com/amdegroot/ssd.pytorch

  • 因为我们是对训练好的pytorch模型进行转换,所以我们可以只要测试的代码就行了,所以我们可以直接对SSD文件夹主目录下的test.py进行修改就行了,修改前先进行备份(以后涉及到修改的文件都可以备份),假设命名为convert2libtorch.py,之后我们对该文件进行修改,

  • 修改开始(按照官网的教程来https://pytorch.org/tutorials/advanced/cpp_export.html,主要是annotation方式,实际上很多复杂点的网络好多control flow):

    • 1.  我们主要对convert2libtorch.py 的test_voc函数进行修改,

      • 先添加torch script用到的包,同时可以添加cv2包(因为测试时可能会用到)

      • 其他包比如

        因为我这里没有下载voc,所以里面很多路径是没有的,但是这样运行会出错,所以我这边的文件改成如下

        如果你是下载过数据集并且完整跑过这些代码,那么上面的可以不改

 

to be continued...

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章