在 C++ 中載入 TorchScript 模型¶

建立於：2018 年 9 月 14 日 | 最後更新：2024 年 12 月 02 日 | 最後驗證：2024 年 11 月 05 日

警告

TorchScript 已停止積極開發。

正如其名稱所示，PyTorch 的主要介面是 Python 程式語言。雖然 Python 是一種適用且首選的語言，適用於許多需要動態性和易於迭代的情況，但在許多情況下，Python 的這些特性同樣是不利的。後者經常應用的環境之一是 *production* – 低延遲和嚴格部署要求的領域。對於生產場景，C++ 通常是首選語言，即使只是將其綁定到另一種語言，如 Java、Rust 或 Go。以下段落將概述 PyTorch 提供的路徑，從現有的 Python 模型到可以完全從 C++ *載入*和*執行*的序列化表示，而無需依賴 Python。

步驟 1：將您的 PyTorch 模型轉換為 Torch Script¶

PyTorch 模型從 Python 到 C++ 的旅程由Torch Script啟用，Torch Script 是 PyTorch 模型的一種表示形式，可以被 Torch Script 編譯器理解、編譯和序列化。如果您從以 vanilla “eager” API 撰寫的現有 PyTorch 模型開始，您必須首先將您的模型轉換為 Torch Script。在下面討論的最常見情況下，這只需要很少的努力。如果您已經有一個 Torch Script 模組，您可以跳到本教學的下一節。

有兩種方法可以將 PyTorch 模型轉換為 Torch Script。第一種稱為 *tracing*，一種透過使用範例輸入評估模型一次並記錄這些輸入流經模型的結構來捕獲模型結構的機制。這適用於對控制流程的使用有限的模型。第二種方法是在您的模型中新增顯式註釋，通知 Torch Script 編譯器它可以直接解析和編譯您的模型程式碼，但須遵守 Torch Script 語言施加的約束。

提示

您可以在官方Torch Script 參考中找到這兩種方法的完整文件，以及有關使用哪種方法的更多指導。

透過 Tracing 轉換為 Torch Script¶

若要透過 tracing 將 PyTorch 模型轉換為 Torch Script，您必須將模型的實例連同範例輸入傳遞給 torch.jit.trace 函式。這將產生一個 torch.jit.ScriptModule 物件，其中模型的 tracing 評估嵌入在模組的 forward 方法中

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

現在可以像評估常規 PyTorch 模組一樣評估 traced 的 ScriptModule

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=<SliceBackward>)

透過註釋轉換為 Torch Script¶

在某些情況下，例如您的模型採用特定形式的控制流程，您可能希望直接在 Torch Script 中編寫您的模型並相應地對您的模型進行註釋。例如，假設您有以下 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

因為此模組的 forward 方法使用依賴於輸入的控制流程，所以它不適合 tracing。相反，我們可以將其轉換為 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 + input
        return output

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)

如果您需要在您的 nn.Module 中排除一些方法，因為它們使用了 TorchScript 尚不支援的 Python 功能，您可以使用 @torch.jit.ignore 對它們進行註釋

sm 是 ScriptModule 的一個實例，已準備好進行序列化。

步驟 2：將您的 Script 模組序列化到檔案¶

一旦你手上有了一個 ScriptModule，無論是透過追蹤 (tracing) 或註解 (annotating) PyTorch 模型而來，你就可以將它序列化到檔案中。之後，你就可以在 C++ 中從這個檔案載入模組，並在不依賴 Python 的情況下執行它。假設我們要序列化先前在追蹤範例中展示的 ResNet18 模型。要執行序列化，只需在模組上呼叫 save 並傳入檔名即可。

traced_script_module.save("traced_resnet_model.pt")

這會在你的工作目錄中產生一個 traced_resnet_model.pt 檔案。如果你也想序列化 sm，請呼叫 sm.save("my_module_model.pt")。現在我們已經正式離開 Python 的領域，準備跨入 C++ 的世界。

步驟 3：在 C++ 中載入你的 Script Module¶

要在 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 {
    // Deserialize the ScriptModule from a file using 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 函式庫中的所有相關 include。我們的應用程式接受序列化的 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 17)

我們建置範例應用程式所需的最後一件事是 LibTorch 發佈版。你始終可以從 PyTorch 網站上的下載頁面取得最新的穩定版本。如果你下載並解壓縮最新的壓縮檔，你應該會收到一個具有以下目錄結構的資料夾：

libtorch/
  bin/
  include/
  lib/
  share/

lib/ 資料夾包含你必須連結的共享函式庫，
include/ 資料夾包含你的程式需要包含的標頭檔案，
share/ 資料夾包含必要的 CMake 設定，以啟用上面簡單的 find_package(Torch) 命令。

提示

在 Windows 上，偵錯 (debug) 和發佈 (release) 版本不具有 ABI 相容性。如果你計劃以偵錯模式建置你的專案，請嘗試 LibTorch 的偵錯版本。此外，請確保你在下面的 cmake --build . 行中指定正確的設定。

最後一步是建置應用程式。為此，假設我們的範例目錄如下所示：

example-app/
  CMakeLists.txt
  example-app.cpp

我們現在可以執行以下命令，從 example-app/ 資料夾中建置應用程式：

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release

其中 /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 的路徑提供給產生的 example-app 二進制檔案，我們應該會得到友好的 "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++ 中執行 Script Module¶

成功在 C++ 中載入我們序列化的 ResNet18 之後，我們現在只需幾行程式碼即可執行它！讓我們將這些行新增到我們的 C++ 應用程式的 main() 函式中：

// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// Execute the model and turn its output into a tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

前兩行設定我們模型的輸入。我們建立一個 torch::jit::IValue (一個型別抹除的值型別，script::Module 方法接受和回傳) 的向量，並新增一個輸入。要建立輸入張量，我們使用 torch::ones()，它等同於 C++ API 中的 torch.ones。然後，我們執行 script::Module 的 forward 方法，將我們建立的輸入向量傳遞給它。作為回傳，我們得到一個新的 IValue，我們透過呼叫 toTensor() 將其轉換為張量。

提示

要了解更多關於 torch::ones 等函式以及一般的 PyTorch C++ API，請參閱其文件：https://pytorch.dev.org.tw/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=<SliceBackward>)

看起來很吻合！

提示

要將您的模型移動到 GPU 記憶體，您可以寫入 model.to(at::kCUDA);。請確保模型的輸入也存在於 CUDA 記憶體中，方法是呼叫 tensor.to(at::kCUDA)，這將會傳回一個新的 CUDA 記憶體中的 Tensor。

步驟 5：取得協助和探索 API¶

希望本教學能讓您對 PyTorch 模型從 Python 到 C++ 的路徑有基本的理解。透過本教學所描述的概念，您應該能夠從一個原始的、“eager”的 PyTorch 模型，到 Python 中的編譯後的 ScriptModule，到磁碟上的序列化檔案，並 - 為了完成這個循環 - 到 C++ 中可執行的 script::Module。

當然，還有很多概念我們沒有涵蓋到。例如，您可能會想要使用 C++ 或 CUDA 中實現的自定義運算符來擴展您的 ScriptModule，並在純 C++ 生產環境中載入的 ScriptModule 中執行此自定義運算符。好消息是：這是可行的，並且得到了很好的支援！現在，您可以探索這個資料夾中的範例，我們將很快推出一個教學。在此期間，以下連結可能普遍有幫助：

Torch Script 參考文檔：https://pytorch.dev.org.tw/docs/master/jit.html
PyTorch C++ API 文件：https://pytorch.dev.org.tw/cppdocs/
PyTorch Python API 文件：https://pytorch.dev.org.tw/docs/

一如既往，如果您遇到任何問題或有疑問，可以使用我們的論壇或 GitHub issue 來與我們聯繫。