使用PyTorch导出JIT模型：C++ API与libtorch实战

最新推荐文章于 2024-09-26 22:19:53 发布

lazycatlove

最新推荐文章于 2024-09-26 22:19:53 发布

阅读量631

点赞数 4

文章标签： pytorch c++ 人工智能

本文链接：https://blog.youkuaiyun.com/lazycatlove/article/details/140723906

版权

PyTorch导出JIT模型并用C++ API libtorch调用

本文将介绍如何将一个 PyTorch 模型导出为 JIT 模型并用 PyTorch 的 C++API libtorch运行这个模型。

Step1：导出模型

首先我们进行第一步，用 Python API 来导出模型，由于本文的重点是在后面的部署阶段，因此，模型的训练就不进行了，直接对 torchvision 中自带的 ResNet50 进行导出。在实际应用中，大家可以对自己训练好的模型进行导出。

# export_jit_model.py
import torch
import torchvision.models as models

model = models.resnet50(pretrained=True)
model.eval()

example_input = torch.rand(1, 3, 224, 224)

jit_model = torch.jit.trace(model, example_input)
torch.jit.save(jit_model, 'resnet50_jit.pth')

导出 JIT 模型的方式有两种：trace 和 script。

我们采用
torch.jit.trace
的方式来导出 JIT 模型，这种方式会根据一个输入将模型跑一遍，然后记录下执行过程。这种方式的问题在于对于有分支判断的模型不能很好的应对，因为一个输入不能覆盖到所有的分支。但是在我们 ResNet50 模型中不会遇到分支判断，因此这里是合适的。关于两种导出 JIT 模型的方式各自优劣不是本文的中断，以后会再写一篇来分析。

在我们的工程目录
demo
下运行上面的
export_jit_model.py
，会得到一个 JIT 模型件：
resnet50_jit.pth
。