今天我們將會跟隨著 PyTorch 的官方腳步,先用 TorchScript 編譯一個模型,能夠在 C++ 的環境下執行。
昨天我們已經說過,可以用 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 的輸出可以說是和未編譯過的模組非常相似。
不同於一般呼叫,torch.save
將訓練結果儲存為 python 的 pickle file,呼叫該編譯過模組 TracedModule save
方法完成。
traced_script_module.save("traced_resnet_model.pt")
下面是一個由官方提供的一個命令列 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
我們在第三部的時候,已經可將在 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 大探索。