iT邦幫忙

第 11 屆 iT 邦幫忙鐵人賽

DAY 9
1
AI & Data

深度學習裡的冰與火之歌 : Tensorflow vs PyTorch系列 第 9

Day 9: TorchScript 實例操作

今天我們將會跟隨著 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 進來的物件所屬的類別。在這裡,我們使用 inspectgetmro 模組級函式,不僅可檢視物件所屬的類別,也可檢視繼承關係。使用方法就是將欲查詢的類別傳入,而回傳的是一個 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 指令。但,首先,你需要以下的路徑:

  1. 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 大探索。


上一篇
Day 8: Tensorflow 在 2.0 很好,那麼 PyTorch 呢?
下一篇
Day 10: 利用 numpy 和 scipy 客製化 PyTorch 模型
系列文
深度學習裡的冰與火之歌 : Tensorflow vs PyTorch31

尚未有邦友留言

立即登入留言