libtorch部署yolov5

一、环境

  • Ubuntu 16.04
  • CUDA 10.2
  • OpenCV 3.4.12
  • LibTorch 1.7.0 (gpu)
  • Pytorch 1.7.0

二、得到yolov5 的权重模型

https://github.com/ultralytics/yolov5/blob/master/models/export.py
我使用的是3.0版本
用export.py生成的pt文件既可,但是默认是使用CPU的,使用GPU需要修改。

提示:用的coremltools==4.0版本,之前我有升级到4.1版本运行export.py会报错,用回4.0

用GPU修改如下:

#img = torch.zeros((opt.batch_size, 3, *opt.img_size))  # image size(1,3,320,192) iDetection
img = torch.zeros((opt.batch_size, 3, *opt.img_size)).to(device='cuda')

#model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float()
 model = torch.load(opt.weights, map_location=torch.device('cuda'))['model'].float()
 
#model.model[-1].export = True 
 model.model[-1].export = False

在终端运行:

python models/export.py --weights yolov5s.pt --img 640 --batch 1  

三、编译运行libtorch-yolov5

git 工程 https://github.com/yasenh/libtorch-yolov5/

cd /path/to/libtorch-yolo5

wget https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.7.0.zip

unzip libtorch-cxx11-abi-shared-with-deps-1.6.0.zip

mkdir build && cd build

cmake .. && make

各个libtorch版本:https://blog.csdn.net/lxx4610/article/details/105806017/

把yolov5工程中生成的pt文件复制到libtorch-yolov5的weights文件夹

编译成功后在终端运行:

./libtorch-yolov5 --source ../images/bus.jpg --weights ../weights/yolov5s.torchscript.pt --gpu --view-img

由于原来我用的opencv 版本是3.4.2
报错:核心转储

terminate called after throwing an instance of ‘cv::Exception’
what(): OpenCV(3.4.2) /tmp/build/80754af9/opencv-suite_1535558553474/work/modules/highgui/src/window.cpp:615: error: (-2:Unspecified error) The function is not implemented. Rebuild the library with Windows, GTK+ 2.x or Carbon support. If you are on Ubuntu or Debian, install libgtk2.0-dev and pkg-config, then re-run cmake or configure script in function ‘cvNamedWindow’

四、opencv源码编译

重新下载opencv源码编译
https://opencv.org/releases/
这次我下载了高版本4.3.0
这个版本可能是高了,也报错。
最后还是改成了3.4.12

unzip opencv-3.4.12.zip
cd opencv--3.4.12

mkdir build && cd build
cmake -D CMAKE_BUILD_TYPE=RELEASE -D OPENCV_GENERATE_PKGCONFIG=ON -D CMAKE_INSTALL_PREFIX=/usr/local/opencv3.4.12 ..
make -j4
sudo make install

编译安装过程很顺利
libtorch部署yolov5_第1张图片
libtorch部署yolov5_第2张图片

这时需要重新编译libtorch-yolov5工程:
修改CMakeLists.txt 文件

#find_package(OpenCV REQUIRED)
#修改成

set(OpenCV_DIR "/usr/local/opencv3.4.12/") 
find_package(OpenCV 3.4.12 REQUIRED)

libtorch部署yolov5_第3张图片
libtorch部署yolov5_第4张图片

五、运行结果

编译过后在终端中运行显示结果:

./libtorch-yolov5 --source ../images/bus.jpg --weights ../weights/yolov5s.torchscript.pt --gpu --view-img

你可能感兴趣的:(深度学习)