libtorch初体验

Posted 2021-02-03 riddick

环境 Ubuntu -18.04.1， opencv3.4.0 ,  python 3.6,  cmake 3.5.0, pytorch 1.0。

pytorch官网下载对应版本：https://pytorch.org/。

libtorch库的安装参考：https://pytorch.org/cppdocs/installing.html。

example代码下载：https://github.com/iamhankai/cpp-pytorch。

 

具体流程和example代码中的一样，先在tracing.py中 将pytorch模型转换为C++可读的script模型，并保存为.pt文件，两条语句就够了：

traced_script_module = torch.jit.trace(model, example)；

traced_script_module.save("model.pt")；

然后在example-app.cpp中加载.pt模型，并使用opencv读取图像数据，做预处理，然后前向预测得到结果。

需要在注意的一点是，example-app.cpp中的一个头文件   #include <torch/Tensor.h>    在libtorch库中是没有的，需要删掉。

最终对dog.png图预测结果如下图：