環境 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圖預測結果如下圖: