多平台轻量级PyTorch模型推理框架MsnhNet

Msnhnet

一款轻量的用于推理pytorch模型的框架,该框架受darknet启发.

msnhnetviewer.png

Ps: 目前本框架在持续开发中,后续将支持GPU版本.也欢迎有志之士一块维护本框架. 欢迎Star.

git地址 https://github.com/msnh2012/Msnhnet

目前测试过的操作系统 (你可以自己测试下mac os)

目前测试过的CPU (其它U出现BUG可以在issue提出)

特点

• 由纯c++编写完成, 不依赖其它第三方计算库, 同时也支持OpenBlas库.
• 该框架自带一套类似netron的可视化工具, 可在编写Msnhnet文件时清晰掌握网络结构.

自己开发的可视本框架的模型

• 支持目前主流的操作系统,Windows, linux(Ubuntu测试), mac os(未测试).
• 支持目前主流的cpu芯片, Intel X86,AMD(未测试) 和ARM.
• 持续开发中...(也就晚上有时间... (╮（╯＿╰）╭))

Yolo测试

• Win10 MSVC 2017 I7-10700F (Linux平台大约比Win快10~20%,当然不是所有网络) nettimeyolov3465msyolov3_tiny75msyolov4600ms
• 树莓派4B nettimeyolov3_tiny2s
• Jeston NX nettimeyolov3_tiny1s

其他已测试的网络(欢迎测试新的模型)

• lenet5
• lenet5_bn
• alexnet
• vgg16
• vgg16_bn
• resnet18
• resnet34
• resnet50
• resnet101
• resnet152
• darknet53
• googLenet
• mobilenetv2
• yolov3
• yolov3_spp
• yolov3_tiny
• yolov4
• 预训练模型下载 链接：https://pan.baidu.com/s/1WElMhBhaN5EnPJnD8S1P3w 提取码：1hlm

依赖

• OpenCV4 https://github.com/opencv/opencv
• yaml-cpp https://github.com/jbeder/yaml-cpp
• Qt5 (可选. 编译Msnhnet viewer时使用) http://download.qt.io/archive/qt/

如何编译

• CMake要求: CMake 3.10+
• 编译选项

cmake_option.png

• ps. 你可以在CMakeLists.txt:43设置最大OMP占用线程, 默认是OMP使用所有CPU线程.(需取消勾选OMP_MAX_THREAD)
• Windows平台编译(MSVC)
  • 1.使用CMake编译安装OpenCV4和Yaml-cpp.
  • 2.在环境变量中添加"OpenCV_DIR"和"yaml-cpp_DIR"变量,并设置变量内容为每个库的CMake目录.
  • 3.下载安装Qt5.
  • 4.把Qt5的bin目录添加环境变量Path.
  • 5.最后使用CMake GUI工具配置Msnhnet然后使用Visual Studio编译安装.
• Linux(Ubuntu)

  sudo apt-get install qt5-default      #optional
  sudo apt-get install libqt5svg5-dev   #optional
  sudo apt-get install libopencv-dev
  
  # build yaml-cpp
  git clone https://github.com/jbeder/yaml-cpp.git
  cd yaml-cpp
  mdir build 
  cd build 
  cmake ..
  make -j4
  sudo make install 
  
  #config 
  sudo echo /usr/local/lib > /etc/ld.so.conf/usrlib.conf
  sudo ldconfig
  
  # build Msnhnet
  git clone https://github.com/msnh2012/Msnhnet.git
  
  cd Msnhnet/build
  cmake -DCMAKE_BUILD_TYPE=Release ..  
  make -j4
  sudo make install
  
  vim ~/.bashrc # Last line add: export PATH=/usr/local/bin:$PATH

测试

• 1.首先下载预训练模型并解压. eg.D:/models.
• 2.打开终端或CMD, cd到安装目录. eg. D:/Msnhnet/bin
• 3.测试 yolov3 "yolov3 D:/models".
• 4.测试 yolov3tiny_video "yolov3tiny_video D:/models".
• 5.测试 classify "classify D:/models".

dog.png

使用MsnhnetViewer查看网络结构

• 1.打开终端或CMD, cd到安装目录 eg. D:/Msnhnet/bin
• 2.运行 "MsnhnetViewer"

viewer.jpg

• Ps. "ResBlock Res2Block AddBlock ConcatBlock"为网络层的封装, 你可以双击该块查看详细内容.

各个封装的结构

ResBlock

Res2Block

Res2Block

AddBlock

Add Block

ConcatBlock

Concat Block

如何转换模型

• 1.使用Pytorch加载你的模型, 并输出网络结构.

import torchvision.models as models
import torch
from torchsummary import summary 

md = models.resnet18(pretrained = True)
md.to("cpu")
md.eval()

print(md, file = open("net.txt", "a"))

summary(md, (3, 224, 224),device='cpu')

• 2.根据网络结构, 编写Msnhnet文件, 跟darknet和caffe类似.
• 3.导出msnhBin (具体细节请查看"pytorch2msnhbin/pytorch2msnhbin.py")

val = []
dd = 0
for name in md.state_dict():
        if "num_batches_tracked" not in name:
                c = md.state_dict()[name].data.flatten().numpy().tolist()
                dd = dd + len(c)
                print(name, ":", len(c))
                val.extend(c)

with open("alexnet.msnhbin","wb") as f:
    for i in val :
        f.write(pack('f',i))

关于训练使用pytorch版本的各个模型训练即可,训练完成后,按照上一步,构建Msnhnet即可. 如yolov3/4. 使用U版: https://github.com/ultralytics/yolov3

推荐语

该框架的作者是中科院知名大佬msnh2012业余时间独立开发的。从框架的立意来看，作者是对各个平台以及各个框架的计算过程有深刻认识的，另外作者也会持续更新支持更多的主流网络例如YOLOV5。并且也在底层加速发力，目前已经支持部分操作的AVX指令集以及Neon指令集优化。