VS配置LibTorch环境并调用CUDA运行torchscript模型

最新推荐文章于 2024-05-22 13:42:19 发布

\lambda

最新推荐文章于 2024-05-22 13:42:19 发布

阅读量2.5k

点赞数 2

分类专栏：深度学习文章标签： c++ opencv pytorch libtorch

本文链接：https://blog.csdn.net/augurlee/article/details/107281309

版权

深度学习专栏收录该内容

21 篇文章 6 订阅

订阅专栏

文章目录

导出 TorchScript 模型
环境配置
示例代码
添加CUDA支持
总结和参考

本文环境： win10 + vs2019 + opencv 4
最新环境：LibTorch：libtorch-win-shared-with-deps-1.6.0+cu101

导出 TorchScript 模型

有两种方式，可以自己试一下，而且很多开源模型已经提供了转换脚本，可以参考

torch.jit.trace
torch.jit.script

还是很简单的：

    # Input
    img = torch.zeros((batch_size, 3, img_size[0], img_size[1])) 
    # Load PyTorch model
    model = torch.load("module.pt", map_location=torch.device('cpu'))['model'].float()
    model.eval()
    model.model[-1].export = True  # set Detect() layer export=True
    y = model(img)  
    # TorchScript export
    try:
        print('\nStarting TorchScript export with torch %s...' % torch.__version__)
        ts = torch.jit.trace(model, img)
        ts.save("module.torchscript")
        print('TorchScript export success, saved as %s' % f)
    except Exception as e:
        print('TorchScript export failure: %s' % e)

环境配置

配置OpenCV属性表

添加目录：
在这里插入图片描述
链接器中加入：

配置LibTorch属性表

下载release版本 850.25M：官网。配置LibTorch，和OpenCV一个套路，属性表我上传了

我下载了一夜，分享出来方便大家 -.-
最新版本 1.5.1 2020/07
天翼网盘：libtorch-win-shared-with-deps-1.5.1.zip （访问码：4sns）

添加【两个】包含目录:
在这里插入图片描述

库目录还是lib
在这里插入图片描述

在输入链接器的时候，由于lib文件比较多，可以通过脚本命令快速得到文件名，然后复制进去就可以了，非常好用。

dir /b *.lib > a.txt

也可以直接复制？

asmjit.lib
c10.lib
c10_cuda.lib
caffe2_detectron_ops_gpu.lib
caffe2_module_test_dynamic.lib
caffe2_nvrtc.lib
clog.lib
cpuinfo.lib
dnnl.lib
fbgemm.lib
libprotobuf-lite.lib
libprotobuf.lib
libprotoc.lib
mkldnn.lib
torch.lib
torch_cpu.lib
torch_cuda.lib

添加环境变量

在这里插入图片描述

示例代码

网上很多代码是测试版本的，接口已经改了。这里是最新的 2020/07

#include "opencv2/opencv.hpp"
#include "torch/script.h"

using namespace std;

torch::jit::script::Module load_model(string model_path)
{
    torch::jit::script::Module module;
    try {
        module = torch::jit::load(model_path);
    }
    catch (const c10::Error & e) {
        std::cerr << "error loading the model\n";
        exit(-1);
    }
    std::cout << "ok\n";
    return module;
}

int main() {
    // 加载模型
    auto module = load_model("modules/yolov5s_dmf.torchscript");

    // 读入图像
    cv::Mat img = cv::imread("test.jpg");
    assert(not img.empty());

    // 转换数据格式
    cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
    cv::Mat img_float;
    img.convertTo(img_float, CV_32F, 1.0 / 255);
    cv::resize(img_float, img_float, cv::Size(640, 640));
    auto img_tensor = torch::from_blob(img_float.data, { 1, 3, 640, 640 }, torch::kFloat32);
    std::vector<torch::jit::IValue> inputs;
    inputs.emplace_back(img_tensor);
    //inputs.push_back(torch::ones({ 1, 3, 448, 640 }));

    // 预测
    at::Tensor output = module.forward(inputs).toTensor();
    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

    return 0;
}

添加CUDA支持

经过上述配置，其实运行的还是cpu环境，是调不起来CUDA的。可以自行测试：

torch::DeviceType device_type;
if (torch::cuda::is_available()) {
    device_type = torch::kCUDA;
    cout << "enable CUDA ...";
} else {
    device_type = torch::kCPU;
}
device = torch::Device(device_type);
this->model.to(device);
this->model.eval();
std::cout << "done.\n";