导出 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)

设置LibTorch环境

下载release版本 850.25M：官网

我下载了一夜，分享出来方便大家 -.-
最新版本 1.5.1 2020/07
天翼网盘：libtorch-win-shared-with-deps-1.5.1.zip （访问码：4sns）
配置环境变量
解压后，添加到path: 你的路径\libtorch\lib
配置OpenCV
配置LibTorch
和OpenCV一个套路，属性表我上传了

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

示例代码

网上很多代码是测试版本的，接口已经改了。这里是最新的 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;
}

总结和参考

网上的许多文章已经过期，LibTorch已经发布了稳定的1.5+版本，还是建议参看英文的官方教程。

在C++ OpenCV 环境中部署PyTorch模型

文章目录

导出 TorchScript 模型

设置LibTorch环境

示例代码

总结和参考