↑ 點(diǎn)擊藍(lán)字?關(guān)注極市平臺(tái)

作者丨火星少女@知乎

來源丨h(huán)ttps://zhuanlan.zhihu.com/p/146453159

編輯丨極市平臺(tái)

極市導(dǎo)讀

本文主要講解如何將pytorch的模型部署到c++平臺(tái)上的模型流程，按順序分為四大塊詳細(xì)說明了模型轉(zhuǎn)換、保存序列化模型、C ++中加載序列化的PyTorch模型以及執(zhí)行Script Module。>>加入極市CV技術(shù)交流群，走在計(jì)算機(jī)視覺的最前沿

最近因?yàn)楣ぷ餍枰裵ytorch的模型部署到c++平臺(tái)上，基本過程主要參照官網(wǎng)的教學(xué)示例，期間發(fā)現(xiàn)了不少坑，特此記錄。

1.模型轉(zhuǎn)換

libtorch不依賴于python，python訓(xùn)練的模型，需要轉(zhuǎn)換為script model才能由libtorch加載，并進(jìn)行推理。在這一步官網(wǎng)提供了兩種方法：

方法一：Tracing

這種方法操作比較簡單，只需要給模型一組輸入，走一遍推理網(wǎng)絡(luò)，然后由torch.ji.trace記錄一下路徑上的信息并保存即可。示例如下：

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)

缺點(diǎn)是如果模型中存在控制流比如if-else語句，一組輸入只能遍歷一個(gè)分支，這種情況下就沒辦法完整的把模型信息記錄下來。

方法二：Scripting

直接在Torch腳本中編寫模型并相應(yīng)地注釋模型，通過torch.jit.script編譯模塊，將其轉(zhuǎn)換為ScriptModule。示例如下：

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

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

forward方法會(huì)被默認(rèn)編譯，forward中被調(diào)用的方法也會(huì)按照被調(diào)用的順序被編譯

如果想要編譯一個(gè)forward以外且未被forward調(diào)用的方法，可以添加[email protected].

如果想要方法不被編譯，可使用

@torch.jit.ignore（https://pytorch.org/docs/master/generated/torch.jit.ignore.html#torch.jit.ignore）?

或者?@torch.jit.unused（https://pytorch.org/docs/master/generated/torch.jit.unused.html#torch.jit.unused）

# Same behavior as pre-PyTorch 1.2
@torch.jit.script
def some_fn():
    return 2

# Marks a function as ignored, if nothing
# ever calls it then this has no effect
@torch.jit.ignore
def some_fn2():
    return 2

# As with ignore, if nothing calls it then it has no effect.
# If it is called in script it is replaced with an exception.
@torch.jit.unused
def some_fn3():
  import pdb; pdb.set_trace()
  return 4

# Doesn't do anything, this function is already
# the main entry point
@torch.jit.export
def some_fn4():
    return 2

在這一步遇到好多坑，主要原因可歸為一下兩點(diǎn)

1. 不支持的操作

TorchScript支持的操作是python的子集，大部分torch中用到的操作都可以找到對(duì)應(yīng)實(shí)現(xiàn)，但也存在一些尷尬的不支持操作，詳細(xì)列表可見https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported，下面列一些我自己遇到的操作：

1）參數(shù)/返回值不支持可變個(gè)數(shù)，例如

def __init__(self, **kwargs):

或者

if output_flag == 0:
    return reshape_logits
else:
    loss = self.loss(reshape_logits, term_mask, labels_id)
    return reshape_logits, loss

2)各種iteration操作

eg1.

layers = [int(a) for a in layers]

報(bào)錯(cuò)torch.jit.frontend.UnsupportedNodeError: ListComp aren’t supported

可以改成：

for k in range(len(layers)):
    layers[k] = int(layers[k])

eg2.

seq_iter = enumerate(scores)
try:
    _, inivalues = seq_iter.__next__()
except:
    _, inivalues = seq_iter.next()

eg3.

line = next(infile)

3)不支持的語句

eg1. 不支持continue

torch.jit.frontend.UnsupportedNodeError: continue statements aren’t supported

eg2. 不支持try-catch

torch.jit.frontend.UnsupportedNodeError: try blocks aren’t supported

eg3. 不支持with語句

4）其他常見op/module

eg1. torch.autograd.Variable

解決：使用torch.ones/torch.randn等初始化+.float()/.long()等指定數(shù)據(jù)類型。

eg2. torch.Tensor/torch.LongTensor etc.

解決：同上

eg3. requires_grad參數(shù)只在torch.tensor中支持，torch.ones/torch.zeros等不可用

eg4. tensor.numpy()

eg5. tensor.bool()

解決：tensor.bool()用tensor>0代替

eg6. self.seg_emb(seg_fea_ids).to(embeds.device)

解決：需要轉(zhuǎn)gpu的地方顯示調(diào)用.cuda()

總之一句話：除了原生python和pytorch以外的庫，比如numpy什么的能不用就不用，盡量用pytorch的各種API。

2. 指定數(shù)據(jù)類型

1）屬性，大部分的成員數(shù)據(jù)類型可以根據(jù)值來推斷，空的列表/字典則需要預(yù)先指定

from typing import Dict

class MyModule(torch.nn.Module):
    my_dict: Dict[str, int]

    def __init__(self):
        super(MyModule, self).__init__()
        # This type cannot be inferred and must be specified
        self.my_dict = {}

        # The attribute type here is inferred to be `int`
        self.my_int = 20

    def forward(self):
        pass

m = torch.jit.script(MyModule())

2）常量，使用Final關(guān)鍵字

try:
    from typing_extensions import Final
except:
    # If you don't have `typing_extensions` installed, you can use a
    # polyfill from `torch.jit`.
    from torch.jit import Final

class MyModule(torch.nn.Module):

    my_constant: Final[int]

    def __init__(self):
        super(MyModule, self).__init__()
        self.my_constant = 2

    def forward(self):
        pass

m = torch.jit.script(MyModule())

3）變量。默認(rèn)是tensor類型且不可變，所以非tensor類型必須要指明

def forward(self, batch_size:int, seq_len:int, use_cuda:bool):

方法三：Tracing and Scriptin混合

一種是在trace模型中調(diào)用script，適合模型中只有一小部分需要用到控制流的情況，使用實(shí)例如下：

import torch

@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r


def bar(x, y, z):
    return foo(x, y) + z

traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))

另一種情況是在script module中用tracing生成子模塊，對(duì)于一些存在script module不支持的python feature的layer，就可以把相關(guān)layer封裝起來，用trace記錄相關(guān)layer流，其他layer不用修改。使用示例如下：

import torch
import torchvision

class MyScriptModule(torch.nn.Module):
    def __init__(self):
        super(MyScriptModule, self).__init__()
        self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
                                        .resize_(1, 3, 1, 1))
        self.resnet = torch.jit.trace(torchvision.models.resnet18(),
                                      torch.rand(1, 3, 224, 224))

    def forward(self, input):
        return self.resnet(input - self.means)

my_script_module = torch.jit.script(MyScriptModule())

2.保存序列化模型

如果上一步的坑都踩完，那么模型保存就非常簡單了，只需要調(diào)用save并傳遞一個(gè)文件名即可，需要注意的是如果想要在gpu上訓(xùn)練模型，在cpu上做inference，一定要在模型save之前轉(zhuǎn)化，再就是記得調(diào)用model.eval(),形如

gpu_model.eval()
cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pth")

traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pth")

3.C++ load訓(xùn)練好的模型

要在C ++中加載序列化的PyTorch模型，必須依賴于PyTorch C ++ API（也稱為LibTorch）。libtorch的安裝非常簡單，只需要在pytorch官網(wǎng)（https://pytorch.org/）下載對(duì)應(yīng)版本，解壓即可。會(huì)得到一個(gè)結(jié)構(gòu)如下的文件夾。

libtorch/ bin/ include/ lib/ share/

然后就可以構(gòu)建應(yīng)用程序了，一個(gè)簡單的示例目錄結(jié)構(gòu)如下：

example-app/ CMakeLists.txt example-app.cpp

example-app.cpp和CMakeLists.txt的示例代碼分別如下：

#include  // One-stop header.
#include #include 
int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app \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;
  }

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

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

至此，就可以運(yùn)行以下命令從example-app/文件夾中構(gòu)建應(yīng)用程序啦：

mkdir build cd build cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. cmake --build . --config Release

其中/path/to/libtorch是之前下載后的libtorch文件夾所在的路徑。這一步如果順利能夠看到編譯完成100%的提示，下一步運(yùn)行編譯生成的可執(zhí)行文件，會(huì)看到“ok”的輸出，可喜可賀！

4. 執(zhí)行Script Module

終于到最后一步啦！下面只需要按照構(gòu)建輸入傳給模型，執(zhí)行forward就可以得到輸出啦。一個(gè)簡單的示例如下：

// Create a vector of inputs.
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';

前兩行創(chuàng)建一個(gè)torch::jit::IValue的向量，并添加單個(gè)輸入. 使用torch::ones()創(chuàng)建輸入張量，等效于C ++ API中的torch.ones。然后，運(yùn)行script::Module的forward方法，通過調(diào)用toTensor()將返回的IValue值轉(zhuǎn)換為張量。C++對(duì)torch的各種操作還是比較友好的，通過torch::或者后加_的方法都可以找到對(duì)應(yīng)實(shí)現(xiàn)，例如

torch::tensor(input_list[j]).to(at::kLong).resize_({batch, 128}).clone()
//torch::tensor對(duì)應(yīng)pytorch的torch.tensor; at::kLong對(duì)應(yīng)torch.int64;resize_對(duì)應(yīng)resize

最后check一下確保c++端的輸出和pytorch是一致的就大功告成啦~

踩了無數(shù)坑，薅掉了無數(shù)頭發(fā)，很多東西也是自己一點(diǎn)點(diǎn)摸索的，如果有錯(cuò)誤歡迎指正！

參考資料：

PyTorch C++ API - PyTorch master document

Torch Script - PyTorch master documentation

文章地址：

https://pytorch.org/cppdocs/

https://pytorch.org/tutorials/advanced/cpp_export.html

推薦閱讀

pytorch的C++ extension寫法
PyTorch擴(kuò)展自定義PyThon/C++(CUDA)算子的若干方法總結(jié)
PyTorch模型量化工具學(xué)習(xí)

添加極市小助手微信（ID : cvmart2），備注：姓名-學(xué)校/公司-研究方向-城市（如：小極-北大-目標(biāo)檢測(cè)-深圳），即可申請(qǐng)加入極市目標(biāo)檢測(cè)/圖像分割/工業(yè)檢測(cè)/人臉/醫(yī)學(xué)影像/3D/SLAM/自動(dòng)駕駛/超分辨率/姿態(tài)估計(jì)/ReID/GAN/圖像增強(qiáng)/OCR/視頻理解等技術(shù)交流群：每月大咖直播分享、真實(shí)項(xiàng)目需求對(duì)接、求職內(nèi)推、算法競賽、干貨資訊匯總、與?10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企視覺開發(fā)者互動(dòng)交流~

△長按添加極市小助手

△長按關(guān)注極市平臺(tái)，獲取最新CV干貨

覺得有用麻煩給個(gè)在看啦~??

在C++平臺(tái)上部署PyTorch模型流程+踩坑實(shí)錄

1.模型轉(zhuǎn)換

2.保存序列化模型

3.C++ load訓(xùn)練好的模型

4. 執(zhí)行Script Module

參考資料：