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

2020-07-10 18:09:33 浏览数 (1)

Msnhnet

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

msnhnetviewer.png

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

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

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

windows

linux

mac os

checked

x

gpu

x

x

x

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

Intel i7

树莓派4B

Jeton NX(cpu)

checked

特点

  • 由纯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)
代码语言:javascript复制
  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加载你的模型, 并输出网络结构.
代码语言:javascript复制
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")
代码语言:javascript复制
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指令集优化。

0 人点赞