最近有個項目需要使用C++調用訓練好的模型。剛好pytorch1.0版本的發布,加入了對C++的支持,准備試一試pytorch對C++的支持怎么樣。這里是官方文檔和教程。
https://pytorch.org/docs/master/jit.htmlhttps://pytorch.org/tutorials/advanced/cpp_export.html
總的來說,現在可以用python版的pytorch快速實現和訓練,使用相應的API導出模型供C++版的pytorch讀取,給C++版本相應輸入會生成和python版本一樣的預測結果。
開發環境
- VS2015(VS2017親測也能通過)
- win10
- cmake>=3.0
轉換模型
pytorch的C++版本用的是Torch Script,官方給了兩種將pytorch模型轉成Torch Script的方法。
第一種方法,Tracing:
這種方法比較簡單,不需要添加代碼到模型中。只需要傳一個輸入給torch.jit.trace函數,讓它輸出一次,然后save。
import Image import torch import torchvision.models as models from torchvision import transforms as transform model_resnet = models.resnet50() #model_resnet.load_state_dict(torch.load("resnet_Epoch_4_Top1_99.75845336914062.pkl")) model_resnet.eval() image = Image.open("your image path").convert('RGB') transforms = transform.Compose([ transform.Resize((224,224)), transform.ToTensor(), transform.Normalize(mean=[0.5]*3, std=[0.5]*3) ]) input = centre_crop_val(image) input = input.unsqueeze(0) traced_script_module_resnet = torch.jit.trace(model_resnet, input) output = traced_script_module_resnet(input) #print(output) traced_script_module_resnet.save("model_resnet_jit.pt")
使用什么做輸出都無所謂,為了方便比較python版和C++版是否輸出一樣,建議使用一個樣本來測試下,不然給對方使用的時候發現結果不一樣就尷尬了(逃。需要和訓練的size以及channel保持一致,同時要保證用於測試的樣本和用於訓練的樣本的transform要一致,不然輸出也不一樣 。使用torch.rand或者torch.ones也是可行的,不會影響已經訓練好的模型權重。
#使用torch.rand
input = torch.ones(1, 3, 224, 224)
traced_script_module_resnet = torch.jit.trace(model_resnet, input)
output = traced_script_module_resnet(input)
#print(output)
traced_script_module_resnet.save("model_resnet_jit.pt")
第二種方法,Annotation:
第二種適合有控制流的模型,比如你的forward方法中有if/else語句,可能就需要使用這種方法。比如用官方的例子做展示:
import torch
class MyModule(torch.nn.Module):
def __init__(self, N, M):
super(MyModule, self).__init__()
self.weight = torch.nn.Parameter(torch.rand(N, M))
def forward(self, input):
if input.sum() > 0:
output = self.weight.mv(input)
else:
output = self.weight + input
return output
對於這種模型,可以在forward方法前加一個修飾器@torch.jit.script_method。
import torch
class MyModule(torch.jit.ScriptModule):
def __init__(self, N, M):
super(MyModule, self).__init__()
self.weight = torch.nn.Parameter(torch.rand(N, M))
@torch.jit.script_method
def forward(self, input):
if bool(input.sum() > 0):
output = self.weight.mv(input)
else:
output = self.weight + input
return output
my_script_module = MyModule(2, 3)
my_script_module.save("your_model.pt")
不管哪種方法得到的model.pt(也就是Torch Script),就可以使用C++調用它了。
准備工作
確定有>=3.0版本的cmake和比較高的vs版本。cmake下載。
在pytorch官網下載對應的LibTorch。有GPU版CP官網下載對應的LibTorch。有GPU版CPU版、有DEBUG和RELEASE版。
然后解壓。
有include有lib,跟其他庫結構差不多。
VS配置
官方和其他很多都是用的cmake,其實vs也能用。新建一個空項目,然后和VS配置opencv一樣,把LibTorch的include和lib添加到“包含目錄”和“庫目錄”中就行,還需要在鏈接器中加入:
torch.lib
c10.lib
caffe2.lib
一般來說3個就足夠,以防萬一可以把所有lib都加上:
c10.lib
caffe2.lib
caffe2_detectron_ops.lib
caffe2_module_test_dynamic.lib
clog.lib
cpuinfo.lib
foxi_dummy.lib
foxi_loader.lib
libprotobuf.lib
libprotobuf-lite.lib
libprotoc.lib
onnx.lib
onnx_proto.lib
onnxifi_dummy.lib
onnxifi_loader.lib
torch.lib
還有兩個地方需要修改:
第一項:屬性->C/C++ ->常規->SDL檢查->否。
第二項:屬性->C/C++ ->語言->符號模式->否。
編寫C++代碼
新建一個example.cpp,選幾張測試的圖片,用opencv讀入然后轉成tensor。訓練網絡的時候Tensor的shape是N x C x H x W,所以還需要把opencv轉成的tensor(H x W x C)用permute轉換一下,然后unsqueeze添加一維變成N x C x H x W。同時要保證測試樣本和訓練樣本有一樣的transform。
#include <torch/script.h> // One-stop header. #include <opencv2/opencv.hpp> #include <iostream> #include <memory> //https://pytorch.org/tutorials/advanced/cpp_export.html std::string image_path = "your image folder path"; int main(int argc, const char* argv[]) { // Deserialize the ScriptModule from a file using torch::jit::load(). std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("your model path"); assert(module != nullptr); std::cout << "ok\n"; //輸入圖像 auto image = cv::imread(image_path +"/"+ "your image name",cv::ImreadModes::IMREAD_IMREAD_COLOR); cv::Mat image_transfomed; cv::resize(image, image_transfomed, cv::Size(70, 70)); // 轉換為Tensor torch::Tensor tensor_image = torch::from_blob(image_transfomed.data, {image_transfomed.rows, image_transfomed.cols,3},torch::kByte); tensor_image = tensor_image.permute({2,0,1}); tensor_image = tensor_image.toType(torch::kFloat); tensor_image = tensor_image.div(255); tensor_image = tensor_image.unsqueeze(0); // 網絡前向計算 at::Tensor output = module->forward({tensor_image}).toTensor(); //std::cout << "output:" << output << std::endl; auto prediction = output.argmax(1); std::cout << "prediction:" << prediction << std::endl; int maxk = 3; auto top3 = std::get<1>(output.topk(maxk, 1, true, true