在C++中加载TorchScript模型
在C++中加載TorchScript模型
本教程已更新為可與PyTorch 1.2一起使用
顧名思義,PyTorch的主要接口是Python編程語言。盡管Python是合適于許多需要動態性和易于迭代的場景,并且是首選的語言,但同樣的,在 許多情況下,Python的這些屬性恰恰是不利的。后者通常適用的一種環境是要求生產-低延遲和嚴格部署。對于生產場景,即使只將C ++綁定到Java, Rust或Go之類的另一種語言中,它也是經常選擇的語言。以下各段將概述PyTorch提供的從現有Python模型到可以完全從C ++加載和執行的序 列化表示形式的路徑,而無需依賴Python。
步驟1:將PyTorch模型轉換為Torch腳本
PyTorch模型從Python到C ++的旅程由Torch Script啟動,Torch Script是PyTorch模型的一種表示形式,可以由Torch Script編譯器理解, 編譯和序列化。如果是從使用vanilla“eager” API編寫的現有PyTorch模型開始的,則必須首先將模型轉換為Torch腳本。在最常見的情況 下(如下所述),這只需要花費很少的功夫。如果已經有了Torch腳本模塊,則可以跳到本教程的下一部分。
有兩種將PyTorch模型轉換為Torch腳本的方法。第一種稱為跟蹤,一種機制,其中通過使用示例輸入對模型的結構進行一次評估,并記錄這些 輸入在模型中的流量,從而捕獲模型的結構。這適用于有限使用控制流的模型。第二種方法是在模型中添加顯式批注,以告知Torch Script編 譯器可以根據Torch Script語言施加的約束直接解析和編譯模型代碼。
提示:可以在官方Torch腳本參考中找到有關這兩種方法的完整文檔,以及使用方法的進一步指導。
方法1:通過跟蹤轉換為Torch腳本
要將PyTorch模型通過跟蹤轉換為Torch腳本,必須將模型的實例以及示例輸入傳遞給torch.jit.trace函數。這將產生一個torch.jit.ScriptModule 對象,該對象的模型評估痕跡將嵌入模塊的forward方法中:
import torch
import torchvision
模型的一個實例.
model = torchvision.models.resnet18()
通常會提供給模型的forward()方法的示例輸入。
example = torch.rand(1, 3, 224, 224)
使用torch.jit.trace來通過跟蹤生成torch.jit.ScriptModule
traced_script_module = torch.jit.trace(model, example)
現在可以對跟蹤的ScriptModule進行評估,使其與常規PyTorch模塊相同:
In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=)
方法2:通過注釋轉換為Torch腳本
在某些情況下,例如,如果模型采用特定形式的控制流,則可能需要直接在Torch腳本中編寫模型并相應地注釋模型。例如,假設具有以下 vanilla Pytorch模型:
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 + inputreturn output
因為此模塊的前向方法使用取決于輸入的控制流,所以它不適合跟蹤。相反,可以將其轉換為ScriptModule。為了將模塊轉換為 ScriptModule,需要使用torch.jit.script編譯模塊,如下所示:
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 + inputreturn output
my_module = MyModule(10,20)
sm = torch.jit.script(my_module)
如果需要在nn.Module中排除某些方法,因為它們使用了TorchScript尚不支持的Python功能,則可以使用@torch.jit.ignore對其進行注釋
my_module是ScriptModule的實例,可以序列化。
步驟2:將腳本模塊序列化為文件
一旦有了ScriptModule(通過跟蹤或注釋PyTorch模型),就可以將其序列化為文件了。稍后,將可以使用C ++從此文件加載模塊并執行它, 而無需依賴Python。假設要序列化先前在跟蹤示例中顯示的ResNet18模型。要執行此序列化,只需在模塊上調用save并傳遞一個文件名即可:
traced_script_module.save(“traced_resnet_model.pt”)
這將在工作目錄中生成traced_resnet_model.pt文件。如果還想序列化my_module,請調用my_module.save(“my_module_model.pt”) 現在已經正式離開Python領域,并準備跨入C ++領域。
步驟3:在C ++中加載腳本模塊
要在C ++中加載序列化的PyTorch模型,應用程序必須依賴于PyTorch C ++ API(也稱為LibTorch)。LibTorch發行版包含共享庫,頭文件 和CMake構建配置文件的集合。雖然CMake不是依賴LibTorch的要求,但它是推薦的方法,并且將來會得到很好的支持。 對于本教程,將 使用CMake和LibTorch構建一個最小的C ++應用程序,該應用程序簡單地加載并執行序列化的PyTorch模型。
最小的C ++應用程序
讓從討論加載模塊的代碼開始。以下將已經做:
include <torch/script.h> // 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 {
// 使用以下命令從文件中反序列化腳本模塊: 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”;
}
<torch/script.h>標頭包含運行示例所需的LibTorch庫中的所有相關包含。的應用程序接受序列化的PyTorch ScriptModule的文件路徑 作為其唯一的命令行參數,然后使用torch::jit::load()函數繼續對該模塊進行反序列化,該函數將此文件路徑作為輸入。作為返回, 收到一個Torch::jit::script::Module對象。將稍后討論如何執行它。
取決于LibTorch和構建應用程序
假設將以上代碼存儲在名為example-app.cpp的文件中。最小的CMakeLists.txt可能看起來很簡單:
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 11)
建立示例應用程序的最后一件事是LibTorch發行版。可以隨時從PyTorch網站的下載頁面上獲取最新的穩定版本。 如果下載并解壓縮最新的歸檔文件,則應收到具有以下目錄結構的文件夾:
libtorch/
bin/
include/
lib/
share/
? lib/ 文件夾包含必須鏈接的共享庫,
? include/ 文件夾包含程序需要包含的頭文件,
? share/ 文件夾包含必要的CMake配置,以啟用上面的簡單find_package(Torch)命令。
提示;在Windows上,調試和發行版本不兼容ABI。 如果打算以調試模式構建項目,請嘗試使用LibTorch的調試版本。
最后一步是構建應用程序。為此,假定示例目錄的布局如下:
example-app/
CMakeLists.txt
example-app.cpp
現在,可以運行以下命令從example-app/文件夾中構建應用程序:
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch …
make
/path/to/libtorch應該是解壓縮的LibTorch發行版的完整路徑。如果一切順利,它將看起來像這樣:
root@4b5a67132e81:/example-app# mkdir build
root@4b5a67132e81:/example-app# cd build
root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch …
– The C compiler identification is GNU 5.4.0
– The CXX compiler identification is GNU 5.4.0
– Check for working C compiler: /usr/bin/cc
– Check for working C compiler: /usr/bin/cc – works
– Detecting C compiler ABI info
– Detecting C compiler ABI info - done
– Detecting C compile features
– Detecting C compile features - done
– Check for working CXX compiler: /usr/bin/c++
– Check for working CXX compiler: /usr/bin/c++ – works
– Detecting CXX compiler ABI info
– Detecting CXX compiler ABI info - done
– Detecting CXX compile features
– Detecting CXX compile features - done
– Looking for pthread.h
– Looking for pthread.h - found
– Looking for pthread_create
– Looking for pthread_create - not found
– Looking for pthread_create in pthreads
– Looking for pthread_create in pthreads - not found
– Looking for pthread_create in pthread
– Looking for pthread_create in pthread - found
– Found Threads: TRUE
– Configuring done
– Generating done
– Build files have been written to: /example-app/build
root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app
如果提供了之前創建的到示例應用程序二進制文件的跟蹤ResNet18模型traced_resnet_model.pt的路徑,則應該以友好的“ ok” 作為獎勵。 請注意,如果嘗試使用my_module_model.pt運行此示例,則會收到一條錯誤消息,提示輸入的形狀不兼容。my_module_model.pt 需要1D而不是4D。
root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt
ok
步驟4:在C ++中執行腳本模塊
成功用C ++加載了序列化的ResNet18之后,現在只需執行幾行代碼即可!讓將這些行添加到C ++應用程序的main()函數中:
// 創建輸入向量
std::vectortorch::jit::IValue inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));
// 執行模型并將輸出轉化為張量
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/dim=/1, /start=/0, /end=/5) << ‘\n’;
前兩行設置了模型的輸入。創建一個torch::jit::IValue的向量(類型為type-erased的值Script::Module方法接受并返回), 并添加單個輸入。要創建輸入張量,使用torch::ones(),等效于C ++ API中的torch.ones。然后,運行script::Module的 forward方法,并向其傳遞創建的輸入向量。作為回報,得到一個新的IValue,通過調用toTensor()將其轉換為張量。
提示:要總體上了解有關torch::ones和PyTorch C ++ API之類的功能的更多信息,請參閱其文檔,網址為https://pytorch.org/cppdocs。 PyTorch C ++ API提供了與Python API幾乎相同的功能奇偶校驗,使可以像在Python中一樣進一步操縱和處理張量。
在最后一行中,打印輸出的前五個條目。由于在本教程前面的部分中,向Python中的模型提供了相同的輸入,因此理想情況下, 應該看到相同的輸出。讓通過重新編譯的應用程序并以相同的序列化模型運行它來進行嘗試:
root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app
root@4b5a67132e81:/example-app/build# ./example-app traced_resnet_model.pt
-0.2698 -0.0381 0.4023 -0.3010 -0.0448
[ Variable[CPUFloatType]{1,5} ]
作為參考,Python以前的輸出為:
tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=)
看來匹配得很好!
提示:要將模型移至GPU內存,可以編寫model.to(at::kCUDA);。通過調用tensor.to(at::kCUDA),確保模型的輸入也位于CUDA內存中, 這將在CUDA內存中返回新的張量。
步驟5:獲取幫助并探索API
本教程有望使對PyTorch模型從Python到C ++的路徑有一個大致的了解。使用本教程中描述的概念,應該能夠從vanilla, “eager” PyTorch模型, 到Python中的已編譯ScriptModule,再到磁盤上的序列化文件,以及–結束循環–到可執行腳本: C ++中的模塊。
當然,有許多沒有介紹的概念。例如,可能會發現自己想要使用以C ++或CUDA實現的自定義運算符擴展ScriptModule,并在加載到純 C ++生產環境中的ScriptModule中執行此自定義運算符。好消息是:這是可能的,并且得到了很好的支持!現在,可以瀏覽此文件夾中的示例, 將很快提供一個教程。 目前,以下鏈接通常可能會有所幫助:
? Torch Script參考:https://pytorch.org/docs/master/jit.html
? PyTorch C ++ API文檔:https://pytorch.org/cppdocs/
? PyTorch Python API文檔:https://pytorch.org/docs/
與往常一樣,如果遇到任何問題或疑問,可以使用的論壇或GitHub issues進行聯系。
總結
以上是生活随笔為你收集整理的在C++中加载TorchScript模型的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 通过带Flask的REST API在Py
- 下一篇: TensorFlow分布式(多GPU和多