Day 9: TorchScript 實例操作

11th鐵人賽

renewang

2019-09-24 23:45:06

19168 瀏覽

分享至

今天我們將會跟隨著 PyTorch 的官方腳步，先用 TorchScript 編譯一個模型，能夠在 C++ 的環境下執行。

第一步: 將你的 PyTorch 模型轉換成 Torch Script

昨天我們已經說過，可以用 Torch Script 中的兩種編譯模式：其一是 tracing compiler，另外則是 script compiler。在此步中，我們會使用 tracing module 作靜態轉換，再用 ScriptModule 作動態轉換。今天我們來嘗試 load 進一個使用 PyTorch python API 寫成的複雜模型，然後再將這個模型使用 TorchScript 編譯。

這裡的例子，所用的 torch 和 torchvision 版本。

import torch
print(torch.__version__)
# => 1.2.0
import torchvision
print(torchvision.__version__)
# => 0.4.0

現在，我們 load 進一個 ResNet，關於這個 implementation 的細節，可以閱讀這篇文章。

import inspect

model = torchvision.models.resnet18() # An instance of your model.
print(inspect.getmro(model.__class__))
#=> (<class 'torchvision.models.resnet.ResNet'>, <class 'torch.nn.modules.module.Module'>, <class 'object'>)

另外我們也使用 inspect 模組，來檢視使用 torchvision.models load 進來的物件所屬的類別。在這裡，我們使用 inspect 的 getmro 模組級函式，不僅可檢視物件所屬的類別，也可檢視繼承關係。使用方法就是將欲查詢的類別傳入，而回傳的是一個 tuple。在 tuple 內的元素，每一個都是該物體在繼承階層中所屬的一個類別，由近至遠。
另外，python 允許多重繼承，所以 getmro 模組級函式也具備 method-resolution-order 來解析多重繼承間的關係由於，除非該類別有實作的特別的 meta-class，每一個元素都應該是類別。
所以，我們檢視輸出，可以看到 ResNet 類別，本身也是一個 torch.nn.Modules，而最後傳回的就是所有物件的 base class，也是 python object 類別。

import numpy as np
example = torch.rand(1, 3, 224, 224) # 提供模型輸入，可以讓模型的 forward 方法使用
traced_script_module = torch.jit.trace(model, example) # 使用 torch.jit.trace 產生一個 torch.jit.ScriptModule 
output1 = model(example)
output2 = traced_script_module(example)
np.testing.assert_allclose(output1.detach().numpy(), output2.detach().numpy()) 

#=> TracedModule[ResNet](
#  (conv1): TracedModule[Conv2d]()
#  (bn1): TracedModule[BatchNorm2d]()
#  (relu): TracedModule[ReLU]()
#  (maxpool): TracedModule[MaxPool2d]()
#  (layer1): TracedModule[Sequential](
#    (0): TracedModule[BasicBlock](
#...
# )

我們可以看到經由 torch.jit.trace 編譯完成的 module 除了都被替換成 TracedModule 外，我們仍能透過 python 呼叫，並傳入相同的輸入來 evaluate 模型的計算。我們可以看到在使用 torch.jit.trace，必須要傳入與呼叫 nn.Modules 類別的 forward 方法，相同的輸入，已完成正向傳播的評估。但，我們也藉著 numpy.testing.assert_allclose 來檢查編譯過的模組和未編譯過的模組行為上是否相似。
我們可以從結果，知道編譯過的 TracedModule 的輸出可以說是和未編譯過的模組非常相似。

第二步: 序列化你已轉換完成的 Script Module 到檔案

不同於一般呼叫，torch.save 將訓練結果儲存為 python 的 pickle file，呼叫該編譯過模組 TracedModule save方法完成。

traced_script_module.save("traced_resnet_model.pt")

第三步: 從檔案載入你已轉換完成的 Script Module 到 C++ 執行環境內

下面是一個由官方提供的一個命令列 C++ application，主要接受一個 filename 的引述，透過 torch script 的 c++ 函式庫，將在第二步寫入的 TracedModule load 進這個 C++ 執行環境。由於我們需要 parse 寫入的 TracedModule 格式，所以我們需要 PyTorch C++ API, 又被稱為 LibTorch。

//example-app.cpp
#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 {
    // 使用 torch::jit::load() 反序列化並讀入 TorchScript 的輸出檔案
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

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

接著，我們必須把這個 C++ 的原始檔編譯成執行檔。在 PyTorch 的官方文件中是介紹 cmake 自動編譯系統。 cmake 是一個可跨平台，可運作在不同的 generator 的自動編譯系統。如在 Unix-based 的架構中，Unix Makefiles 是預設選擇的 generator，其他的編譯架構選擇還包括了 Ninja 或 xcode 等。這個自動編譯系統也可以與 IDE 結合，如 Eclipse CDT4，或 Sublime Text 4。
不過，今天我們不用 cmake 來做編譯，而是用簡單的命令列中的 one-liner 指令。但，首先，你需要以下的路徑：

path_to_libtorch: 通常位於安裝 PyTorch 的位置，如 <path to virtual environment>/lib/python3.7/site-packages/torch/。則這個 torch 檔案夾則是 libtorch 所在位置。 libtorch 有下列幾個重要的檔案結構：

torch/
  bin/
  include/<header files>
         torch/script.h # 在此應該能發現 include 的 torch/script.h 
  lib/<library files>
         libc10.dylib    #c10 Error 需要
         libtorch.dylib  #torch::jit 方法需要
  share/

根據我的 MacOSX 系統，library 是編譯成 dylib 檔，若是在 Unix-based 則有可能是編譯成 *.so 檔。
2. path_to_std：因為我所用的編譯系統是 llvm + clang，該系統有自行實踐的 std c++ library，所以必須要給定 std c++ header 的位置，且透過 -std 指定為 c++11 語法。

最後這個 one-liner 編譯指令為

clang++ example-app.cpp -o example-app -std=c++11 -I<path_to_std> -I<path_to_libtorch>/include/ -L<path_to_libtorch>lib/ -ltorch -lc10 -Wl,-rpath <path_to_libtorch>lib/

執行結果

# 為給定路徑
⨯ ./example-app 
usage: example-app <path-to-exported-script-module>
# 若有給定路徑，則是
⨯ ./example-app traced_resnet_model.pt                   
ok

第四步: 在 C++ 執行 Script Module

我們在第三部的時候，已經可將在 python 中利用 trace 模組編譯的模型 load 進 C++ 執行環境內。這次，我們加入 C++ 輸入，並且讀出輸出，印出並和 python 的列印結果比對。
底下是一小段 C++ 的原始碼，可以加在 main 中，loading module 的原始碼後，依照上面的方法做 one-line 編譯，最後執行 ./example-app traced_resnet_model.pt

// 使用 std::vector 建立輸入輸入
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));
// 執行模型並轉換輸出為一個 torch::Tensor 物件
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

接著我們在命令列執行以編譯好的執行檔，並使用剛才利用 tracing compiler 輸出到硬碟的檔案位置作為執行檔的引數，執行的結果如下：

./example-app traced_resnet_model.pt                                        
 ok
0.9397  0.7853  0.6534  0.8342  0.3897
[ Variable[CPUFloatType]{1,5} ]

接著，讓我們光陰倒回，回到一開始還在 python 的時候，用以編譯的模組將前五個輸出印出：

output = traced_script_module(torch.ones(1, 3, 224, 224))
print(output[0, :5])
#=>tensor([0.9397, 0.7853, 0.6534, 0.8342, 0.3897], grad_fn=<SliceBackward>)

從兩者的輸出可以用目光區別，兩者還真相似。最後用這一個結論，正式完結今天的 TorchScript 大探索。