利用LibTorch调用PyTorch训练好的模型

最新推荐文章于 2024-07-31 17:28:45 发布

TangCode

最新推荐文章于 2024-07-31 17:28:45 发布

阅读量6.3k

点赞数 2

分类专栏：编程语言文章标签： PyTorch LibTorch C++ 模型部署

本文链接：https://blog.csdn.net/qq_17232031/article/details/95105462

版权

本文介绍了如何利用C++结合LibTorch库调用PyTorch训练好的ResNet50模型。首先，从PyTorch官网下载LibTorch并解压。接着，使用PyTorch训练模型并保存模型状态。在C++中，配置CMakeLists以链接LibTorch库和OpenCV，并编写代码读取图像，转换格式，然后通过LibTorch的API加载模型，进行前向传播计算，最终得到与Python相同的结果。

摘要由CSDN通过智能技术生成

PyTorch如今发布到1.1稳定版本，新增的功能让模型部署变得更为地简单，本文记录如何利用C++来调用PyTorch训练好的模型，其实也是利用官方强大的LibTorch库。

LibTorch的安装

虽然说安装，其实就是下载官方的LibTorch包而已，从官方网站中选择PyTorch(1.1)，libtorch，以及cuda的版本，其中会出现下载链接，这里为cuda9.0的链接
https://download.pytorch.org/libtorch/cu90/libtorch-shared-with-deps-latest.zip
下载好找个路径解压。解压完放在那不动！！

PyTorch模型训练

这里我使用了最为简单ResNet50的预训练模型，其中保存跟踪模型的代码如下：

import torch
import torchvision.models as models
from PIL import Image
import numpy as np
image = Image.open("build/airliner.jpg") #图片发在了build文件夹下
image = image.resize((224, 224),Image.ANTIALIAS)
image = np.asarray(image)
image = image / 255
image = torch.Tensor(image).unsqueeze_(dim=0)
image = image.permute((0, 3, 1, 2)).float()

model = models.resnet50(pretrained=True)
model = model.eval()
resnet = torch.jit.trace(model, torch.rand(1,3,224,224))
#