之前一直用Pytorch做的网络训练,最近想看看训练出来的模型在C++上的表现(毕竟稍微注重效率的工程都用的C++写),因此花了点时间研究了一波。 因为PyTorch的Python版本和C++版本使用的文件描述有区别,因此需要先把Python版本训练的模型转换成Torch Script形式(相信这一点不太方便的地方后面会被FB大佬们改掉的)。 转换的代码如下:
import torch
import torchvision
# 读取resnet18模型
model = torchvision.models.resnet18()
# 创建一个输入样例
example = torch.rand(1, 3, 224, 224)
# 初始化模型转换工具
traced_script_module = torch.jit.trace(model, example)
# 转换模型
traced_script_module.save("resnet18.pt")
用上面的代码得到模型文件“resnet18.pt”,准备的工作就做完了。 下面到官网下载程序包
解压完里面是这样的
感觉很亲切,因为基本上所有的库都长这鸟样。下一步也跟所有其他库一样,我们新建一个C++工程,然后把上图中的“include”和“lib”往里面扔,再添加到工程里面的“包含目录”和“库目录”即可。 链接库一般会用到下面三个,添加进去
torch.lib
c10.lib
caffe2.lib
到这步就基本做完了。但编译会出错,需要再设置两项参数: 第一项:属性->C/C++ ->常规->SDL检查->否 第二项:属性->C/C++ ->语言->符号模式->否 搞定! 完成了上面简单的配置,用下面的代码就可以在c++上运行程序啦(编译出来缺少dll,则在之前下的libtorch包里面找)!
#include <torch/script.h>
#include <iostream>
#include <memory>
int main()
{
// 读取网络模型
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("resnet18.pt");
assert(module != nullptr);
std::cout << "读取正常\n";
// 创建输入
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({ 1, 3, 224, 224 }));
// 数据输入网络,获得输出
at::Tensor output = module->forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << std::endl;
}
参考文章:https://blog.csdn.net/weixin_41576121/article/details/85110423