因为我的主业是v-slam所以大部分的代码还是在c++中，所以一些调用网络的地方还是需要利用pytorch c++的接口的。
某乎中的一位答主已经解释得很清楚了link，现在拿mnist lenet做一个实验
把模型转为Torch Script
省略mnist lenet的训练过程（这部分tutorial或者网上已经很多了）
...
net = LeNet()
net = net.cuda(device)
net.load_state_dict(torch.load('./model/net_008.pth')) ## 已经训练好的mnist 模型
print('load net success \n')
example_tensor = torch.rand(1,1,28,28).to(device)
traced_script_model = torch.jit.trace(net,example_tensor)
traced_script_model.save("mnist.pt")
调用Tensor Script
下载libtorch,官网已经提供了下载的链接。但是发现有一些ABI的bug。所以我在link这里下载了libtorch。
从测试集中选取一张图片进行测试，如下：
test.png
使用pytorch c++接口调用上一步生成的mnist.pt
#include "torch/script.h"
#include "torch/torch.h"
#include <iostream>
#include <opencv2/opencv.hpp>
int main(int argc, char **argv) {
if (argc != 3)
return 1;
cv::Mat img = cv::imread(argv[1], cv::IMREAD_GRAYSCALE);
cv::resize(img, img, cv::Size(28, 28));
cv::imshow("number", img);
cv::waitKey(0);
torch::jit::script::Module module = torch::jit::load(argv[2]);
torch::DeviceType device_type;
device_type = torch::kCUDA;
torch::Device device(device_type);
img.convertTo(img, CV_32FC1, 1.f / 255.f, 0);
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
std::vector<int64_t> dims = {1, 1, img.rows, img.cols};
auto img_var = torch::from_blob(img.data, dims, torch::kFloat32).to(device);
inputs.push_back(img_var);
auto output = module.forward(inputs).toTensor();
std::cout << output << std::endl;
return 1;
}
输出的结果为：
7.0965 -6.2461 -2.4858 -6.0416 -0.3752 0.9906 9.3611 -9.5642 2.1789 -4.8047
[ Variable[CUDAFloatType]{1,10} ]
结果正确。
以上就是简单测试了pytorch c++接口。更多的内容还是去官网上查比较详细。https://pytorch.org/cppdocs/
