【libtorch模型部署】VS2017中部署libtorch模型

C++中使用libtorch进行推理

---

前言

用libtorch部署pytorch模型，而不是用tensorrt等工具部署模型的优势在于：pytorch和libtorch同属一个生态，API语句比较接近，并且不会出现某网络层不支持的问题。理论上只要是pytorch训练的模型都能用libtorch部署，因为他们共用相同的c++代码。具体的部署流程如下：
1.模型训练：首先使用PyTorch对模型进行训练，得到一个可以使用的模型。
2.模型保存：训练好的模型需要保存为TorchScript模型，这是一个可以被TorchScript解释器执行的模型。
3.模型加载：使用PyTorch的torch.jit.load函数来加载保存好的TorchScript模型。
4.模型部署：使用LibTorch的API，将加载好的模型进行部署，并进行决策
此外，使用libtorch，我们还可以将PyTorch 模型导出为可执行文件，并在没有Python 解释器的情况下运行这些模型，这使得我们可以将训练好的模型部署到嵌入式设备上。

一、libtorch的安装

libtorch下载链接https://pytorch.org/，需要注意的是根据自己的cuda和cudnn版本下载对应的安装包.!

如何查看cuda和cudnn版本，参考
https://blog.csdn.net/weixin_61995249/article/details/124069914。
本文所配置环境为VS2017+cuda11.6+cudnn8.4.0

二、部署步骤

1.导出libtorch部署需要的模型

libtorch不依赖于python，但python训练的模型，需要转换为script model才能由libtorch加载，并进行推理。在这一步官网提供了两种方法：
Trace 和 Script
Tracing导出模型，这种比较简单，缺点是适用于模型没有分支的情况。且有两种方法，一种是训练过程中直接整个模型保存了，第二种是训练过程中仅保存了模型权重，需要再次构建模型再进行转换。可以参考https://blog.csdn.net/long11350/article/details/125803997。

import torch
import torchvision.models as models
from PIL import Image
import numpy as np
 
device = torch.device('cpu')
model = torch.load("test.pth", map_location=device)
 
model.eval()
# print(model)
 
example = torch.rand(1, 1, 224, 224)
 
traced_script_module = torch.jit.trace(model, example)
 
output = traced_script_module(torch.ones(1, 1, 224, 224))
print(output)
 
# ----------------------------------
traced_script_module.save("test.pt")

如果模型较复杂，forward中有分支，那么就需要用通过torch.jit.script编译模块，将其转换为ScriptModule。参考https://blog.csdn.net/mist99/article/details/118671555。 如下所示：

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.