Pytorch1.0的(C++ & cpu)版本在Windows上的使用

84
1
2021年2月6日 09时23分

之前一直用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

发表评论

后才能评论

评论列表(1条)

  • bug404 2021年2月7日 上午8:38

    老哥继续更啊,我也对c艹的torch很感兴趣