环境
Win10、VS2019、OpenCV4.5.1、libtorch1.8.0+cuda11.1.1
下载OpenCV
下载libtorch和pytorch
注意:下载的版本要和YOLOv5的pytorch的版本一致!!!
解压之后:
配置环境变量
扫描二维码关注公众号,回复:
14640965 查看本文章
VS配置
快捷键alt+F7或右击解决方案->属性
选择VC++目录-包含目录,添加以下路径:
选择VC++目录-库目录,添加以下路径:
选择链接器-输入-附加依赖项,添加以下路径:
选择CIC++常规-SDL检查,修改为"否“:
选择链接器-命令行-其它选项,添加/INCLUDE:?warp_size@cuda@at@@YAHXZ
操作完之后就可以使用torch库
测试代码
#include<iostream>
#include<torch/script.h>
#include <torch/torch.h>
#include<memory>
int main()
{
std::cout << "cuda::is_available():" << torch::cuda::is_available() << std::endl;
system("pause");
}
运行模型
我们开始YOLOv5中的export.py文件,首先我们将训练好的best.pt模型放入
因为文件默认生成CPU版本的,所以我们需要改动两处:
运行结束后我们会得到一个best.torchscript.pt文件,这个就是我们所需要的
代码调配
我们总共需要准备三个文件:类别文件、模型文件、图片
输入三个文件的路径
基本到这就没问题了,有问题的话请留言。完事收工