1. 程式人生 > 程式設計 >在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,224))
In[2]: output[0,:5]
Out[2]: tensor([-0.2698,-0.0381,0.4023,-0.3010,-0.0448],grad_fn=<SliceBackward>)

方法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 + input
  return output

因為此模組的前向方法使用取決於輸入的控制流,所以它不適合跟蹤。相反,我們可以將其轉換為ScriptModule 。為了將模組轉換為ScriptModule ,需要使用torch.jit.script 編譯模組,如下所示:

class MyModule(torch.nn.Module):
 def __init__(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)

如果您需要在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 <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().
 module = torch::jit::load(argv[1]);
 }
 catch (const c10::Error& e) {
 std::cerr << "error loading the model\n";
 return -1;
 }

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

標頭包含執行示例所需的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/

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::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1,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,grad_fn=<SliceBackward>)

看來匹配得很好!

提示:要將模型移至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模型的方法,希望對大家有所幫助!