由于python的易用性，深度學(xué)習(xí)模型多是在python框架下進(jìn)行訓(xùn)練的，如TensorFlow，pytorch等。而由于硬件設(shè)備的限制，有時(shí)候其部署可能需要基于C++/C的平臺(tái)。比如在我們的項(xiàng)目中，語義分割網(wǎng)絡(luò)是在pytorch下訓(xùn)練的，而分割結(jié)果基于應(yīng)用的后處理部分是在C下面實(shí)現(xiàn)的，那怎么才能把這兩種平臺(tái)下的東西結(jié)合起來一起運(yùn)行呢？我想到的方法有以下幾種：

• 將pytorch模型訓(xùn)練好后轉(zhuǎn)成caffe的模型，然后利用caffe的接口在C++下面實(shí)現(xiàn)模型的推理應(yīng)用；

• 將C部分的代碼打包編譯成一個(gè)動(dòng)態(tài)連接庫dll，然后在python框架下調(diào)用該dll實(shí)現(xiàn)c下面的功能；

• 利用pytorch的C++版本LibTorch實(shí)現(xiàn)pytorch模型的調(diào)用。
  本文主要記錄最后一種方法。

LibTorch 的下載及使用

LibTorch 是pytorch的C++版本，在pytorch版本1.0后就有了。在官網(wǎng)通過如下選擇，就可以得到下載鏈接。

image-20200304141339075.png

下載鏈接里有release版本和debug版本，建議兩個(gè)版本都下載，兩者主要是對(duì)應(yīng)的dll和lib不一樣，debug版本還提供了pdb，可以幫助定位錯(cuò)誤位置。將release版本解壓后，得到一個(gè)LibTorch的文件夾，再將debug版本解壓，將其中的lib文件夾改名為lib_debug，同樣放在之前release版本解壓的LibTorch文件夾，這樣就release和debug版本都可以使用了。

下載的LibTorch中提供了cmakelist，在linux平臺(tái)可以利用cmake來使用它。而如果在Windows平臺(tái)利用vs，只需要和一般的第三方庫使用一樣，在對(duì)應(yīng)的工程中添加正確的AdditionalIncludeDirectories，AdditionalLibraryDirectories，AdditionalDependencies等就可以了。我在實(shí)驗(yàn)時(shí)，將lib里邊所有的lib文件都加入到AdditionalDependencies了。程序運(yùn)行的時(shí)候還需要把對(duì)應(yīng)的dll拷貝到exe所在的文件夾。我使用debug時(shí)遇到了一個(gè)編譯錯(cuò)誤，添加preprocessor _SCL_SECURE_NO_WARNINGS就好了。

使用流程

利用LibTorch來調(diào)用pytorch模型的流程大致是這樣的：

1. pytorch訓(xùn)練好模型
2. 將模型序列化并存成pt文件
3. 在C中利用LibTorch的接口進(jìn)行正向推演

pytorch模型序列化

第一步我們就不介紹了，我們從第二步開始。模型的序列化是利用Torch Script來完成的。TorchScript是一種從PyTorch代碼創(chuàng)建可序列化和可優(yōu)化模型的方法。用TorchScript編寫的任何代碼都可以從Python進(jìn)程中保存并加載到?jīng)]有Python依賴關(guān)系的進(jìn)程中。對(duì)于一個(gè)已經(jīng)訓(xùn)練好的pytorch模型，官方提供兩種方法進(jìn)行Torch Script的轉(zhuǎn)換：tracing和annotation。

Tracing

Tracing的方法還是很簡(jiǎn)單的，參見如下示例代碼：

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

Annotation

tracing適用于大多數(shù)網(wǎng)絡(luò)，如果你的網(wǎng)絡(luò)的forward方法中對(duì)input有邏輯判斷，比如input的size為一個(gè)值時(shí)走向一個(gè)分支，而為另一值時(shí)走向另一個(gè)分支，那么只能用annotation進(jìn)行轉(zhuǎn)換。比如如下的網(wǎng)絡(luò)：

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

利用annotation來將上述網(wǎng)絡(luò)模型轉(zhuǎn)成Torch Script可以按如下代碼：

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)

annotation的方法我并沒有測(cè)試，我使用的模型用tracing就已經(jīng)足夠了。

序列化

序列化的意思是指將上述Torch Script描述的模型存成一個(gè)文件。

traced_script_module.save("traced_resnet_model.pt")

C++中的正向推演

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }
    
  // Simple tests of the model
  std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::ones({1, 3, 224, 224}));

  // Execute the model and turn its output into a tensor.
  at::Tensor output = module.forward(inputs).toTensor();

  std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

  std::cout << "ok\n";
}

自己訓(xùn)練的模型的實(shí)際操作

下面以我們訓(xùn)練的語義分割網(wǎng)絡(luò)為例，介紹如何將自己的模型在C++中跑起來。

在實(shí)際的操作中，也是遇到了一些問題的。

GPU及DataParallel的問題

第一個(gè)問題是我們之前的模型訓(xùn)練是在GPU（相信應(yīng)該都是這樣的）中進(jìn)行的，并且使用了DataParallel，在序列化時(shí)，如下代碼是正確的，可以與示例代碼做下比較。

device = torch.device('cuda')
model = get_model(args_in)
model = torch.nn.DataParallel(model, device_ids=[0])
model.load_state_dict(torch.load(args_in.test_model_path))
model.to(device)
# use evaluation mode to ignore dropout, etc
model.eval()

# The tracing input need not to be the same size as the forward case.
example = torch.rand(1, 3, 1080, 1920).to(device)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model.module, example)

traced_script_module.save("traced_model.pt")

對(duì)于GPU訓(xùn)練的模型，需要將模型和tracing用的tensor通過to(device)或者.cuda()轉(zhuǎn)到GPU上，如第5,10行。對(duì)于利用DataParallel訓(xùn)練的模型，需要在trace時(shí)使用model.module，如第13行。

關(guān)于DataParallel多說一句，如果希望正向的時(shí)候不需要像第3行那樣將model再包一層，在訓(xùn)練save model的時(shí)候應(yīng)該按如下

torch.save(model.module.state_dict(), save_path)

這樣存的model就不需要第3行代碼，而且第13行的.module也不需要了。

附上因?yàn)镈ataParallel沒弄對(duì)在pycharm中遇到的錯(cuò)誤

RuntimeError: hasSpecialCase INTERNAL ASSERT FAILED at ..\torch\csrc\jit\passes\alias_analysis.cpp:300, please report a bug to PyTorch. We don't have an op for aten::to but it isn't a special case. (analyzeImpl at ..\torch\csrc\jit\passes\alias_analysis.cpp:300)

網(wǎng)絡(luò)輸出是Tuple的問題

我們的網(wǎng)絡(luò)輸出是一個(gè)tuple而不是一個(gè)tensor，于是在C++調(diào)用的時(shí)候總是crash，用了debug版本的LibTorch，才發(fā)現(xiàn)問題。官方提到LibTorch這種方式需要網(wǎng)絡(luò)的輸出是一個(gè)tuple或者tensor，那如果輸出的是tuple，在C++端代碼應(yīng)該按如下修改

torch::Tensor result = module.forward(input).toTuple()->elements()[0].toTensor();

圖像的前處理

在pytorch模型的訓(xùn)練過程中，我們一般會(huì)對(duì)圖像進(jìn)行一些前處理，比如

transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ])

在LibTorch中，可以這樣做

  tensor_image = tensor_image.toType(torch::kFloat);
  tensor_image = tensor_image.div(255);
  // Normalization
  tensor_image[0][0] = tensor_image[0][0].sub_(0.485).div_(0.229);
  tensor_image[0][1] = tensor_image[0][1].sub_(0.456).div_(0.224);
  tensor_image[0][2] = tensor_image[0][2].sub_(0.406).div_(0.225);

最后貼上我們利用opencv讀視頻，然后對(duì)每一幀運(yùn)行語義分割正向的代碼。

  // module for forward process
  torch::jit::script::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load("traced_model.pt");
  } catch (const c10::Error &e) {
    std::cerr << "error loading the model\n";
  }
  torch::DeviceType device = torch::kCUDA;
  module.to(device);

  // opencv windows
  cv::namedWindow("Test", 0);
  cvMoveWindow("Test", 0, 0);
  
  cv::VideoCapture  t_video_in(videoPath);
  long nbFrames = static_cast<long>(t_video_in.get(CV_CAP_PROP_FRAME_COUNT));

  for (long f = 0; f < nbFrames; f++) {
    cv::Mat image, input;
    t_video_in >> image;
    cv::cvtColor(image, input, CV_BGR2RGB);

    // run semantic segmentation to get label image
    torch::Tensor tensor_image = torch::from_blob(input.data, { 1, input.rows, input.cols, 3 }, torch::kByte);
    tensor_image = tensor_image.permute({ 0, 3, 1, 2 });
    tensor_image = tensor_image.toType(torch::kFloat);
    tensor_image = tensor_image.div(255);
    // Normalization
    tensor_image[0][0] = tensor_image[0][0].sub_(0.485).div_(0.229);
    tensor_image[0][1] = tensor_image[0][1].sub_(0.456).div_(0.224);
    tensor_image[0][2] = tensor_image[0][2].sub_(0.406).div_(0.225);

    tensor_image = tensor_image.to(torch::kCUDA);
    torch::Tensor result = module.forward({ tensor_image }).toTuple()->elements()[0].toTensor();
    torch::Tensor pred = result.argmax(1);
    pred = pred.squeeze();
    pred = pred.to(torch::kU8);
    pred = pred.to(torch::kCPU);

    cv::Mat label(cv::Size(image.cols,image.rows), CV_8U, pred.data_ptr());
    cv::imshow("Test", label);

    cv::waitKey(1);
  }
  t_video_in.release();