使用自訂 C++ 運算子擴充 TorchScript

建立於：2018 年 11 月 28 日 | 最後更新：2024 年 7 月 22 日 | 最後驗證：2024 年 11 月 05 日

警告

此教學已自 PyTorch 2.4 起棄用。請參閱 PyTorch 自訂運算子，以取得有關 PyTorch 自訂運算子的最新指南。

PyTorch 1.0 版本向 PyTorch 引入了一個新的程式設計模型，稱為 TorchScript。TorchScript 是 Python 程式語言的一個子集，可以由 TorchScript 編譯器進行解析、編譯和最佳化。此外，編譯後的 TorchScript 模型可以選擇序列化為磁碟檔案格式，您可以隨後從純 C++（以及 Python）載入和執行它以進行推論。

TorchScript 支援 torch 套件提供的大量運算子子集，讓您可以將許多複雜的模型純粹表示為 PyTorch「標準函式庫」中的一系列張量運算。然而，有時您可能會發現自己需要使用自訂 C++ 或 CUDA 函式來擴充 TorchScript。雖然我們建議您僅在您的想法無法（有效率地）表示為簡單的 Python 函式時才訴諸此選項，但我們確實提供了一個非常友善且簡單的介面，用於使用 ATen（PyTorch 的高效能 C++ 張量函式庫）定義自訂 C++ 和 CUDA 核心。一旦綁定到 TorchScript 中，您可以將這些自訂核心（或「運算」）嵌入到您的 TorchScript 模型中，並在 Python 中以及在其序列化形式中直接在 C++ 中執行它們。

以下段落提供了一個編寫 TorchScript 自訂運算以呼叫 OpenCV（一個用 C++ 編寫的電腦視覺函式庫）的範例。我們將討論如何在 C++ 中使用張量，如何有效地將它們轉換為第三方張量格式（在本例中為 OpenCV Mat），如何向 TorchScript 執行時註冊您的運算子，以及最後如何編譯運算子並在 Python 和 C++ 中使用它。

在 C++ 中實作自訂運算子

在本教學中，我們將公開 warpPerspective 函式（將透視轉換應用於圖像），從 OpenCV 作為自訂運算子到 TorchScript。第一步是在 C++ 中編寫自訂運算子的實作。讓我們將此實作的檔案命名為 op.cpp，並使其如下所示

torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
  // BEGIN image_mat
  cv::Mat image_mat(/*rows=*/image.size(0),
                    /*cols=*/image.size(1),
                    /*type=*/CV_32FC1,
                    /*data=*/image.data_ptr<float>());
  // END image_mat

  // BEGIN warp_mat
  cv::Mat warp_mat(/*rows=*/warp.size(0),
                   /*cols=*/warp.size(1),
                   /*type=*/CV_32FC1,
                   /*data=*/warp.data_ptr<float>());
  // END warp_mat

  // BEGIN output_mat
  cv::Mat output_mat;
  cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{8, 8});
  // END output_mat

  // BEGIN output_tensor
  torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});
  return output.clone();
  // END output_tensor
}

此運算子的程式碼非常簡短。在檔案的頂部，我們包含 OpenCV 標頭檔 opencv2/opencv.hpp，以及 torch/script.h 標頭，它公開了 PyTorch C++ API 中所有必要的內容，我們需要編寫自訂 TorchScript 運算子。我們的函式 warp_perspective 採用兩個參數：一個輸入 image 和我們希望應用於圖像的 warp 轉換矩陣。這些輸入的類型是 torch::Tensor，這是 PyTorch 在 C++ 中的張量類型（也是 Python 中所有張量的基礎類型）。我們的 warp_perspective 函式的傳回類型也將是 torch::Tensor。

提示

有關 ATen（向 PyTorch 提供 Tensor 類別的函式庫）的更多資訊，請參閱 此註釋。此外，本教學描述了如何在 C++ 中分配和初始化新的張量物件（此運算子不需要）。

注意

TorchScript 編譯器理解固定數量的型別。只有這些型別可以用作您自訂運算子的引數。目前這些型別包括：torch::Tensor、torch::Scalar、double、int64_t 以及這些型別的 std::vector。請注意，僅支援 double 而不支援 float，並且僅支援 int64_t 而不支援其他整數型別，例如 int、short 或 long。

在我們的函式內部，我們需要做的第一件事是將 PyTorch tensors 轉換為 OpenCV 矩陣，因為 OpenCV 的 warpPerspective 需要 cv::Mat 物件作為輸入。幸運的是，有一種方法可以做到這一點，而無需複製任何資料。在前幾行，

  cv::Mat image_mat(/*rows=*/image.size(0),
                    /*cols=*/image.size(1),
                    /*type=*/CV_32FC1,
                    /*data=*/image.data_ptr<float>());

我們呼叫 OpenCV Mat 類別的 這個建構子，將我們的 tensor 轉換為 Mat 物件。我們將原始 image tensor 的行數和列數、資料類型（在這個範例中，我們將其固定為 float32），以及指向底層資料的原始指標（float*）傳遞給它。關於 Mat 類別的這個建構子，特別之處在於它不會複製輸入資料。相反，它只會引用此記憶體，以供對 Mat 執行的所有操作使用。如果在 image_mat 上執行原地操作，這將反映在原始 image tensor 中（反之亦然）。這允許我們使用該函式庫的原生矩陣型別呼叫後續的 OpenCV 常式，即使我們實際上將資料儲存在 PyTorch tensor 中。我們重複此程序，將 warp PyTorch tensor 轉換為 warp_mat OpenCV 矩陣。

  cv::Mat warp_mat(/*rows=*/warp.size(0),
                   /*cols=*/warp.size(1),
                   /*type=*/CV_32FC1,
                   /*data=*/warp.data_ptr<float>());

接下來，我們準備好呼叫我們渴望在 TorchScript 中使用的 OpenCV 函式：warpPerspective。為此，我們將 image_mat 和 warp_mat 矩陣以及一個名為 output_mat 的空輸出矩陣傳遞給 OpenCV 函式。我們還指定了我們希望輸出矩陣（影像）的大小 dsize。在這個範例中，它被硬編碼為 8 x 8。

  cv::Mat output_mat;
  cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{8, 8});

我們的自訂運算子實作中的最後一步是將 output_mat 轉換回 PyTorch tensor，以便我們可以在 PyTorch 中進一步使用它。這與我們之前在另一個方向轉換時所做的事情非常相似。在這種情況下，PyTorch 提供了一個 torch::from_blob 方法。在這種情況下，*blob* 旨在表示一些不透明的、平坦的記憶體指標，我們希望將其解釋為 PyTorch tensor。torch::from_blob 的呼叫如下所示：

  torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});
  return output.clone();

我們在 OpenCV Mat 類別上使用 .ptr<float>() 方法，以取得指向底層資料的原始指標（就像之前 PyTorch tensor 的 .data_ptr<float>() 一樣）。我們還指定了 tensor 的輸出形狀，我們將其硬編碼為 8 x 8。torch::from_blob 的輸出然後是一個 torch::Tensor，指向 OpenCV 矩陣擁有的記憶體。

在從我們的運算子實作中傳回此 tensor 之前，我們必須在 tensor 上呼叫 .clone()，以執行底層資料的記憶體複製。這樣做的原因是 torch::from_blob 傳回的 tensor 不擁有其資料。在該時間點，資料仍然由 OpenCV 矩陣擁有。但是，此 OpenCV 矩陣將超出範圍，並在函式結束時被取消配置。如果我們按原樣傳回 output tensor，當我們在函式外部使用它時，它將指向無效記憶體。呼叫 .clone() 會傳回一個新的 tensor，其中包含新 tensor 自己擁有的原始資料的副本。因此，可以安全地傳回給外界。

使用 TorchScript 註冊自訂運算子

現在我們已經在 C++ 中實作了自訂運算子，我們需要使用 TorchScript 執行階段和編譯器註冊它。這將允許 TorchScript 編譯器解析 TorchScript 程式碼中對我們的自訂運算子的引用。如果您曾經使用過 pybind11 函式庫，我們的註冊語法與 pybind11 語法非常相似。要註冊單個函式，我們寫入：

TORCH_LIBRARY(my_ops, m) {
  m.def("warp_perspective", warp_perspective);
}

位於 op.cpp 檔案的頂層某處。TORCH_LIBRARY 巨集會建立一個在程式啟動時會被呼叫的函式。您的程式庫名稱 (my_ops) 作為第一個參數給定 (不應加引號)。第二個參數 (m) 定義了一個 torch::Library 類型的變數，它是註冊運算子的主要介面。方法 Library::def 實際上會建立一個名為 warp_perspective 的運算子，並將其公開給 Python 和 TorchScript。您可以透過多次呼叫 def 來定義任意數量的運算子。

在幕後，def 函式實際上做了相當多的工作：它使用樣板元編程來檢查函式的類型簽名，並將其轉換為一個運算子綱要，該綱要指定了 TorchScript 類型系統中的運算子類型。

建立自訂運算子

現在我們已經在 C++ 中實作了自訂運算子並編寫了其註冊碼，現在是時候將該運算子建置為一個（共享）程式庫，我們可以將其載入到 Python 中進行研究和實驗，或載入到 C++ 中以在無 Python 環境中進行推論。可以使用純 CMake 或 Python 替代方案 (例如 setuptools) 以多種方式建置我們的運算子。為了簡潔起見，以下段落僅討論 CMake 方法。本教學課程的附錄深入探討了其他替代方案。

環境設定

我們需要安裝 PyTorch 和 OpenCV。取得它們最簡單且與平台無關的方式是透過 Conda。

conda install -c pytorch pytorch
conda install opencv

使用 CMake 建置

若要使用 CMake 建置系統將我們的自訂運算子建置到共享程式庫中，我們需要編寫一個簡短的 CMakeLists.txt 檔案並將其與先前的 op.cpp 檔案放在一起。為此，讓我們就一個如下所示的目錄結構達成共識：

warp-perspective/
  op.cpp
  CMakeLists.txt

我們的 CMakeLists.txt 檔案的內容應該如下：

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(warp_perspective)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

# Define our library target
add_library(warp_perspective SHARED op.cpp)
# Enable C++14
target_compile_features(warp_perspective PRIVATE cxx_std_14)
# Link against LibTorch
target_link_libraries(warp_perspective "${TORCH_LIBRARIES}")
# Link against OpenCV
target_link_libraries(warp_perspective opencv_core opencv_imgproc)

現在要建置我們的運算子，我們可以從我們的 warp_perspective 資料夾執行以下命令：

$ mkdir build
$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
-- 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
-- Found torch: /libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /warp_perspective/build
$ make -j
Scanning dependencies of target warp_perspective
[ 50%] Building CXX object CMakeFiles/warp_perspective.dir/op.cpp.o
[100%] Linking CXX shared library libwarp_perspective.so
[100%] Built target warp_perspective

這將在 build 資料夾中放置一個 libwarp_perspective.so 共享程式庫檔案。在上面的 cmake 命令中，我們使用輔助變數 torch.utils.cmake_prefix_path 來方便地告訴我們 PyTorch 安裝的 cmake 檔案在哪裡。

我們將在下面進一步詳細探討如何使用和呼叫我們的運算子，但為了盡早獲得成功感，我們可以嘗試在 Python 中執行以下程式碼：

import torch
torch.ops.load_library("build/libwarp_perspective.so")
print(torch.ops.my_ops.warp_perspective)

如果一切順利，這應該會印出類似以下內容：

<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f618fc6fa50>

這是我們稍後將用來呼叫自訂運算子的 Python 函式。

在 Python 中使用 TorchScript 自訂運算子

一旦我們的自訂運算子被建置到共享程式庫中，我們就可以在 Python 中的 TorchScript 模型中使用這個運算子。這分為兩個部分：首先將運算子載入到 Python 中，然後在 TorchScript 程式碼中使用該運算子。

您已經了解如何將運算子匯入到 Python 中：torch.ops.load_library()。這個函式採用包含自訂運算子的共享程式庫的路徑，並將其載入到目前程序中。載入共享程式庫也會執行 TORCH_LIBRARY 區塊。這會向 TorchScript 編譯器註冊我們的自訂運算子，並允許我們在 TorchScript 程式碼中使用該運算子。

您可以將載入的運算子稱為 torch.ops.<namespace>.<function>，其中 <namespace> 是運算子名稱的命名空間部分，而 <function> 是運算子的函式名稱。對於我們上面編寫的運算子，命名空間是 my_ops，函式名稱是 warp_perspective，這意味著我們的運算子可用作 torch.ops.my_ops.warp_perspective。雖然這個函式可以在腳本或追蹤的 TorchScript 模組中使用，但我們也可以直接在 vanilla eager PyTorch 中使用它，並傳遞常規 PyTorch 張量。

import torch
torch.ops.load_library("build/libwarp_perspective.so")
print(torch.ops.my_ops.warp_perspective(torch.randn(32, 32), torch.rand(3, 3)))

產生

tensor([[0.0000, 0.3218, 0.4611,  ..., 0.4636, 0.4636, 0.4636],
      [0.3746, 0.0978, 0.5005,  ..., 0.4636, 0.4636, 0.4636],
      [0.3245, 0.0169, 0.0000,  ..., 0.4458, 0.4458, 0.4458],
      ...,
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000],
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000],
      [0.1862, 0.1862, 0.1692,  ..., 0.0000, 0.0000, 0.0000]])

注意

幕後發生的是，當您第一次在 Python 中存取 torch.ops.namespace.function 時，TorchScript 編譯器（在 C++ 領域中）會查看是否已註冊函式 namespace::function，如果是，則傳回此函式的 Python 控制代碼，我們可以隨後使用該控制代碼從 Python 呼叫到我們的 C++ 運算子實作中。這是 TorchScript 自訂運算子和 C++ 擴充之間的一個值得注意的區別：C++ 擴充使用 pybind11 手動繫結，而 TorchScript 自訂運算子由 PyTorch 本身即時繫結。Pybind11 在您可以繫結到 Python 的類型和類別方面為您提供更大的靈活性，因此建議用於純 eager 程式碼，但 TorchScript 運算子不支援它。

從這裡開始，您可以像使用 torch 封裝中的其他函式一樣，在腳本或追蹤的程式碼中使用您的自訂運算子。事實上，「標準程式庫」函式 (例如 torch.matmul) 經歷的註冊路徑與自訂運算子基本相同，這使得自訂運算子在 TorchScript 中使用方式和位置方面真正成為一等公民。（但是，一個區別是標準程式庫函式具有自訂編寫的 Python 參數剖析邏輯，該邏輯不同於 torch.ops 參數剖析。）

使用追蹤的自訂運算子

讓我們先將運算子嵌入到追蹤的函式中。回想一下，對於追蹤，我們從一些 vanilla Pytorch 程式碼開始

def compute(x, y, z):
    return x.matmul(y) + torch.relu(z)

然後在其上呼叫 torch.jit.trace。我們進一步將一些範例輸入傳遞給 torch.jit.trace，它會將這些範例輸入轉發到我們的實作中，以記錄當輸入流經它時發生的操作序列。這樣做的結果實際上是 eager PyTorch 程式的「凍結」版本，TorchScript 編譯器可以進一步分析、最佳化和序列化它。

inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(4, 5)]
trace = torch.jit.trace(compute, inputs)
print(trace.graph)

產生

graph(%x : Float(4:8, 8:1),
      %y : Float(8:5, 5:1),
      %z : Float(4:5, 5:1)):
  %3 : Float(4:5, 5:1) = aten::matmul(%x, %y) # test.py:10:0
  %4 : Float(4:5, 5:1) = aten::relu(%z) # test.py:10:0
  %5 : int = prim::Constant[value=1]() # test.py:10:0
  %6 : Float(4:5, 5:1) = aten::add(%3, %4, %5) # test.py:10:0
  return (%6)

現在，令人興奮的啟示是，我們可以簡單地將我們的自訂運算子放入我們的 PyTorch 追蹤中，就像它是 torch.relu 或任何其他 torch 函式一樣。

def compute(x, y, z):
    x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
    return x.matmul(y) + torch.relu(z)

然後像以前一樣追蹤它。

inputs = [torch.randn(4, 8), torch.randn(8, 5), torch.randn(8, 5)]
trace = torch.jit.trace(compute, inputs)
print(trace.graph)

產生

graph(%x.1 : Float(4:8, 8:1),
      %y : Float(8:5, 5:1),
      %z : Float(8:5, 5:1)):
  %3 : int = prim::Constant[value=3]() # test.py:25:0
  %4 : int = prim::Constant[value=6]() # test.py:25:0
  %5 : int = prim::Constant[value=0]() # test.py:25:0
  %6 : Device = prim::Constant[value="cpu"]() # test.py:25:0
  %7 : bool = prim::Constant[value=0]() # test.py:25:0
  %8 : Float(3:3, 3:1) = aten::eye(%3, %4, %5, %6, %7) # test.py:25:0
  %x : Float(8:8, 8:1) = my_ops::warp_perspective(%x.1, %8) # test.py:25:0
  %10 : Float(8:5, 5:1) = aten::matmul(%x, %y) # test.py:26:0
  %11 : Float(8:5, 5:1) = aten::relu(%z) # test.py:26:0
  %12 : int = prim::Constant[value=1]() # test.py:26:0
  %13 : Float(8:5, 5:1) = aten::add(%10, %11, %12) # test.py:26:0
  return (%13)

將 TorchScript 自訂運算子整合到追蹤的 PyTorch 程式碼中就像這樣簡單！

使用腳本的自訂運算子

除了追蹤（tracing）之外，另一種取得 PyTorch 程式的 TorchScript 表示形式的方法是直接用 TorchScript 編寫程式碼。TorchScript 基本上是 Python 語言的一個子集，但有一些限制，使 TorchScript 編譯器更容易推斷程式。您可以通過使用 @torch.jit.script 裝飾自由函數和使用 @torch.jit.script_method 裝飾類別中的方法（類別也必須繼承自 torch.jit.ScriptModule），將常規 PyTorch 程式碼轉換為 TorchScript。有關 TorchScript 註解的更多詳細資訊，請參閱這裡。

使用 TorchScript 而非追蹤的一個特殊原因是，追蹤無法捕捉 PyTorch 程式碼中的控制流程。因此，讓我們考慮這個確實使用控制流程的函數

def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  return x.matmul(y) + z

要將此函數從 vanilla PyTorch 轉換為 TorchScript，我們使用 @torch.jit.script 進行裝飾

@torch.jit.script
def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  return x.matmul(y) + z

這會即時編譯 compute 函數為圖形表示形式，我們可以在 compute.graph 屬性中檢查它

>>> compute.graph
graph(%x : Dynamic
    %y : Dynamic) {
  %14 : int = prim::Constant[value=1]()
  %2 : int = prim::Constant[value=0]()
  %7 : int = prim::Constant[value=42]()
  %z.1 : int = prim::Constant[value=5]()
  %z.2 : int = prim::Constant[value=10]()
  %4 : Dynamic = aten::select(%x, %2, %2)
  %6 : Dynamic = aten::select(%4, %2, %2)
  %8 : Dynamic = aten::eq(%6, %7)
  %9 : bool = prim::TensorToBool(%8)
  %z : int = prim::If(%9)
    block0() {
      -> (%z.1)
    }
    block1() {
      -> (%z.2)
    }
  %13 : Dynamic = aten::matmul(%x, %y)
  %15 : Dynamic = aten::add(%13, %z, %14)
  return (%15);
}

現在，就像之前一樣，我們可以在腳本程式碼中使用自定義運算符，就像使用其他函數一樣

torch.ops.load_library("libwarp_perspective.so")

@torch.jit.script
def compute(x, y):
  if bool(x[0] == 42):
      z = 5
  else:
      z = 10
  x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
  return x.matmul(y) + z

當 TorchScript 編譯器看到對 torch.ops.my_ops.warp_perspective 的引用時，它會找到我們透過 C++ 中的 TORCH_LIBRARY 函數註冊的實現，並將其編譯為其圖形表示形式

>>> compute.graph
graph(%x.1 : Dynamic
    %y : Dynamic) {
    %20 : int = prim::Constant[value=1]()
    %16 : int[] = prim::Constant[value=[0, -1]]()
    %14 : int = prim::Constant[value=6]()
    %2 : int = prim::Constant[value=0]()
    %7 : int = prim::Constant[value=42]()
    %z.1 : int = prim::Constant[value=5]()
    %z.2 : int = prim::Constant[value=10]()
    %13 : int = prim::Constant[value=3]()
    %4 : Dynamic = aten::select(%x.1, %2, %2)
    %6 : Dynamic = aten::select(%4, %2, %2)
    %8 : Dynamic = aten::eq(%6, %7)
    %9 : bool = prim::TensorToBool(%8)
    %z : int = prim::If(%9)
      block0() {
        -> (%z.1)
      }
      block1() {
        -> (%z.2)
      }
    %17 : Dynamic = aten::eye(%13, %14, %2, %16)
    %x : Dynamic = my_ops::warp_perspective(%x.1, %17)
    %19 : Dynamic = aten::matmul(%x, %y)
    %21 : Dynamic = aten::add(%19, %z, %20)
    return (%21);
  }

請特別注意圖形末尾對 my_ops::warp_perspective 的引用。

注意

TorchScript 圖形表示形式仍然可能發生變化。請不要依賴它看起來像這樣。

這就是在 Python 中使用自定義運算符的全部內容。簡而言之，您可以使用 torch.ops.load_library 導入包含您的運算符的庫，並從您的追蹤或腳本化的 TorchScript 程式碼中像調用其他 torch 運算符一樣調用您的自定義運算符。

在 C++ 中使用 TorchScript 自定義運算符

TorchScript 的一個有用的功能是能夠將模型序列化到磁碟上的文件中。該文件可以透過網路傳送、儲存在檔案系統中，或者更重要的是，可以動態反序列化和執行，而無需保留原始程式碼。這在 Python 中是可行的，但在 C++ 中也是可行的。為此，PyTorch 提供了 一個純 C++ API，用於反序列化和執行 TorchScript 模型。如果您還沒有閱讀關於在 C++ 中載入和執行序列化 TorchScript 模型的教程，請閱讀它，接下來的幾個段落將以此為基礎。

簡而言之，即使從檔案反序列化並在 C++ 中運行，自定義運算符也可以像常規 torch 運算符一樣執行。唯一的必要條件是將我們之前構建的自定義運算符共享庫與我們執行模型的 C++ 應用程式連結。在 Python 中，這可以簡單地透過調用 torch.ops.load_library 來實現。在 C++ 中，您需要使用您使用的任何構建系統將共享庫與您的主要應用程式連結。以下範例將使用 CMake 來展示這一點。

注意

從技術上講，您也可以在運行時以與我們在 Python 中相同的方式將共享庫動態載入到您的 C++ 應用程式中。在 Linux 上，您可以使用 dlopen 來執行此操作。在其他平台上存在等效項。

在上面連結的 C++ 執行教程的基礎上，讓我們從一個最小的 C++ 應用程式開始，該應用程式位於一個文件中，即 main.cpp，它與我們的自定義運算符位於不同的資料夾中，該應用程式載入並執行序列化的 TorchScript 模型

#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;
  }

  // Deserialize the ScriptModule from a file using torch::jit::load().
  torch::jit::script::Module module = torch::jit::load(argv[1]);

  std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::randn({4, 8}));
  inputs.push_back(torch::randn({8, 5}));

  torch::Tensor output = module.forward(std::move(inputs)).toTensor();

  std::cout << output << std::endl;
}

以及一個小的 CMakeLists.txt 文件

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(example_app)

find_package(Torch REQUIRED)

add_executable(example_app main.cpp)
target_link_libraries(example_app "${TORCH_LIBRARIES}")
target_compile_features(example_app PRIVATE cxx_range_for)

此時，我們應該能夠構建應用程式

$ mkdir build
$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
-- 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
-- Found torch: /libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /example_app/build
$ make -j
Scanning dependencies of target example_app
[ 50%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o
[100%] Linking CXX executable example_app
[100%] Built target example_app

並且可以在沒有傳遞模型的情況下運行它

$ ./example_app
usage: example_app <path-to-exported-script-module>

接下來，讓我們序列化我們之前編寫的使用自定義運算符的腳本函數

torch.ops.load_library("libwarp_perspective.so")

@torch.jit.script
def compute(x, y):
  if bool(x[0][0] == 42):
      z = 5
  else:
      z = 10
  x = torch.ops.my_ops.warp_perspective(x, torch.eye(3))
  return x.matmul(y) + z

compute.save("example.pt")

最後一行會將腳本函數序列化到一個名為「example.pt」的文件中。如果我們然後將這個序列化的模型傳遞給我們的 C++ 應用程式，我們就可以立即運行它

$ ./example_app example.pt
terminate called after throwing an instance of 'torch::jit::script::ErrorReport'
what():
Schema not found for node. File a bug report.
Node: %16 : Dynamic = my_ops::warp_perspective(%0, %19)

或者可能沒有。可能還不行。當然！我們還沒有將自定義運算符庫與我們的應用程式連結。讓我們現在就這樣做，為了正確地這樣做，讓我們稍微更新一下我們的文件組織，使其看起來像這樣

example_app/
  CMakeLists.txt
  main.cpp
  warp_perspective/
    CMakeLists.txt
    op.cpp

這將允許我們將 warp_perspective 庫 CMake 目標添加為我們的應用程式目標的一個子目錄。example_app 資料夾中的頂層 CMakeLists.txt 應該看起來像這樣

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(example_app)

find_package(Torch REQUIRED)

add_subdirectory(warp_perspective)

add_executable(example_app main.cpp)
target_link_libraries(example_app "${TORCH_LIBRARIES}")
target_link_libraries(example_app -Wl,--no-as-needed warp_perspective)
target_compile_features(example_app PRIVATE cxx_range_for)

這個基本的 CMake 配置看起來與以前非常相似，除了我們將 warp_perspective CMake 構建添加為一個子目錄。一旦其 CMake 程式碼運行，我們就將我們的 example_app 應用程式與 warp_perspective 共享庫連結。

注意

在上面的範例中嵌入了一個關鍵細節：-Wl,--no-as-needed 前綴到 warp_perspective 連結行。這是必需的，因為我們實際上不會在我們的應用程式程式碼中調用 warp_perspective 共享庫中的任何函數。我們只需要 TORCH_LIBRARY 函數運行。不方便的是，這會混淆連結器，並使其認為它可以完全跳過與該庫的連結。在 Linux 上，-Wl,--no-as-needed 標誌強制連結發生（注意：此標誌特定於 Linux！）。還有其他解決方法。最簡單的方法是在運算符庫中定義一些函數，您需要從主應用程式中調用該函數。這可以像在某些標頭中聲明的函數 void init(); 一樣簡單，然後在運算符庫中將其定義為 void init() { }。在主應用程式中調用這個 init() 函數會給連結器一個印象，即這是一個值得連結的庫。不幸的是，這超出了我們的控制範圍，我們寧願讓您知道原因和簡單的解決方法，而不是向您提供一些不透明的巨集來放入您的程式碼中。

現在，由於我們在頂層找到了 Torch 套件，因此 warp_perspective 子目錄中的 CMakeLists.txt 檔案可以稍微縮短一些。 它應該看起來像這樣

find_package(OpenCV REQUIRED)
add_library(warp_perspective SHARED op.cpp)
target_compile_features(warp_perspective PRIVATE cxx_range_for)
target_link_libraries(warp_perspective PRIVATE "${TORCH_LIBRARIES}")
target_link_libraries(warp_perspective PRIVATE opencv_core opencv_photo)

讓我們重新建置我們的範例應用程式，它也會與自定義運算子函式庫連結。 在頂層 example_app 目錄中

$ mkdir build
$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
-- 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
-- Found torch: /libtorch/lib/libtorch.so
-- Configuring done
-- Generating done
-- Build files have been written to: /warp_perspective/example_app/build
$ make -j
Scanning dependencies of target warp_perspective
[ 25%] Building CXX object warp_perspective/CMakeFiles/warp_perspective.dir/op.cpp.o
[ 50%] Linking CXX shared library libwarp_perspective.so
[ 50%] Built target warp_perspective
Scanning dependencies of target example_app
[ 75%] Building CXX object CMakeFiles/example_app.dir/main.cpp.o
[100%] Linking CXX executable example_app
[100%] Built target example_app

如果我們現在執行 example_app 二進位檔案，並將我們的序列化模型交給它，我們應該會得到一個圓滿的結局

$ ./example_app example.pt
11.4125   5.8262   9.5345   8.6111  12.3997
 7.4683  13.5969   9.0850  11.0698   9.4008
 7.4597  15.0926  12.5727   8.9319   9.0666
 9.4834  11.1747   9.0162  10.9521   8.6269
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
10.0000  10.0000  10.0000  10.0000  10.0000
[ Variable[CPUFloatType]{8,5} ]

成功！ 您現在已準備好進行推論。

結論

本教學課程逐步引導您如何在 C++ 中實作自定義 TorchScript 運算子，如何將其建置到共享函式庫中，如何使用它在 Python 中定義 TorchScript 模型，以及最後如何將其載入到 C++ 應用程式中以進行推論工作負載。 您現在已準備好使用與第三方 C++ 函式庫介面的 C++ 運算子來擴展您的 TorchScript 模型、編寫自定義的高效能 CUDA 核心，或實作任何其他需要 Python、TorchScript 和 C++ 之間平滑融合的用例。

與往常一樣，如果您遇到任何問題或有任何疑問，您可以使用我們的論壇或 GitHub issues 與我們聯繫。 此外，我們的 常見問題 (FAQ) 頁面 可能會有有用的資訊。

附錄 A：建置自定義運算子的更多方法

「建置自定義運算子」一節說明瞭如何使用 CMake 將自定義運算子建置到共享函式庫中。 本附錄概述了另外兩種編譯方法。 它們都使用 Python 作為編譯過程的「驅動程式」或「介面」。 此外，兩者都重複使用 PyTorch 為 現有的基礎架構 提供的 *C++ 擴充功能*，它們是 TorchScript 自定義運算子的 vanilla (eager) PyTorch 等效項，依賴 pybind11 來將 C++ 中的函數「顯式」綁定到 Python 中。

第一種方法使用 C++ 擴充功能的 方便的即時 (JIT) 編譯介面，在您第一次執行程式碼時在 PyTorch 腳本的背景中編譯您的程式碼。 第二種方法依賴於備受推崇的 setuptools 套件，並涉及編寫一個單獨的 setup.py 檔案。 這允許更進階的配置以及與其他基於 setuptools 的專案整合。 我們將在下面詳細探討這兩種方法。

使用 JIT 編譯建置

PyTorch C++ 擴充工具包提供的 JIT 編譯功能允許將自定義運算子的編譯直接嵌入到您的 Python 程式碼中，例如在您的訓練腳本的頂部。

注意

此處的「JIT 編譯」與 TorchScript 編譯器中發生的 JIT 編譯以優化您的程式無關。 它僅表示您的自定義運算子 C++ 程式碼將在您第一次匯入它時在系統的 /tmp 目錄下的資料夾中編譯，就像您事先自己編譯它一樣。

此 JIT 編譯功能有兩種形式。 第一種形式是，您仍然將運算子實作保留在單獨的檔案中 (op.cpp)，然後使用 torch.utils.cpp_extension.load() 編譯您的擴充功能。 通常，此函數將傳回公開您的 C++ 擴充功能的 Python 模組。 但是，由於我們沒有將我們的自定義運算子編譯到它自己的 Python 模組中，我們只想編譯一個普通的共享函式庫。 幸運的是，torch.utils.cpp_extension.load() 有一個參數 is_python_module，我們可以將其設定為 False 以表明我們只對建置共享函式庫而不是 Python 模組感興趣。 然後，torch.utils.cpp_extension.load() 將編譯並將共享函式庫載入到目前流程中，就像之前的 torch.ops.load_library 一樣

import torch.utils.cpp_extension

torch.utils.cpp_extension.load(
    name="warp_perspective",
    sources=["op.cpp"],
    extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
    is_python_module=False,
    verbose=True
)

print(torch.ops.my_ops.warp_perspective)

這應該大約印出

<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f3e0f840b10>

第二種形式的 JIT 編譯允許您將自定義 TorchScript 運算子的原始碼作為字串傳遞。 為此，請使用 torch.utils.cpp_extension.load_inline

import torch
import torch.utils.cpp_extension

op_source = """
#include <opencv2/opencv.hpp>
#include <torch/script.h>

torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) {
  cv::Mat image_mat(/*rows=*/image.size(0),
                    /*cols=*/image.size(1),
                    /*type=*/CV_32FC1,
                    /*data=*/image.data<float>());
  cv::Mat warp_mat(/*rows=*/warp.size(0),
                   /*cols=*/warp.size(1),
                   /*type=*/CV_32FC1,
                   /*data=*/warp.data<float>());

  cv::Mat output_mat;
  cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{64, 64});

  torch::Tensor output =
    torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{64, 64});
  return output.clone();
}

TORCH_LIBRARY(my_ops, m) {
  m.def("warp_perspective", &warp_perspective);
}
"""

torch.utils.cpp_extension.load_inline(
    name="warp_perspective",
    cpp_sources=op_source,
    extra_ldflags=["-lopencv_core", "-lopencv_imgproc"],
    is_python_module=False,
    verbose=True,
)

print(torch.ops.my_ops.warp_perspective)

當然，最佳實務是僅在您的原始碼相當短時才使用 torch.utils.cpp_extension.load_inline。

請注意，如果您在 Jupyter Notebook 中使用它，則不應多次執行註冊的儲存格，因為每次執行都會註冊一個新的函式庫並重新註冊自定義運算子。 如果您需要重新執行它，請事先重新啟動筆記本的 Python 核心。

使用 Setuptools 建置

完全從 Python 建置我們的自定義運算子的第二種方法是使用 setuptools。 這樣做的好處是 setuptools 具有相當強大且廣泛的介面，用於建置以 C++ 撰寫的 Python 模組。 但是，由於 setuptools 實際上是為建置 Python 模組而不是普通的共享函式庫（它們沒有 Python 從模組中期望的必要進入點）而設計的，因此這條路線可能有點古怪。 也就是說，您所需要的只是一個 setup.py 檔案來代替 CMakeLists.txt，如下所示

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(
    name="warp_perspective",
    ext_modules=[
        CppExtension(
            "warp_perspective",
            ["example_app/warp_perspective/op.cpp"],
            libraries=["opencv_core", "opencv_imgproc"],
        )
    ],
    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)

請注意，我們在底部的 BuildExtension 中啟用了 no_python_abi_suffix 選項。 這指示 setuptools 在產生的共享函式庫的名稱中省略任何 Python-3 特定的 ABI 後綴。 否則，例如在 Python 3.7 上，該函式庫可能會被命名為 warp_perspective.cpython-37m-x86_64-linux-gnu.so，其中 cpython-37m-x86_64-linux-gnu 是 ABI 標籤，但我們實際上只是希望它被命名為 warp_perspective.so

現在，如果我們在終端機中，從 setup.py 所在的資料夾內執行 python setup.py build develop，我們應該會看到類似下面的輸出：

$ python setup.py build develop
running build
running build_ext
building 'warp_perspective' extension
creating build
creating build/temp.linux-x86_64-3.7
gcc -pthread -B /root/local/miniconda/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/torch/csrc/api/include -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/TH -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/THC -I/root/local/miniconda/include/python3.7m -c op.cpp -o build/temp.linux-x86_64-3.7/op.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=warp_perspective -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
creating build/lib.linux-x86_64-3.7
g++ -pthread -shared -B /root/local/miniconda/compiler_compat -L/root/local/miniconda/lib -Wl,-rpath=/root/local/miniconda/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.7/op.o -lopencv_core -lopencv_imgproc -o build/lib.linux-x86_64-3.7/warp_perspective.so
running develop
running egg_info
creating warp_perspective.egg-info
writing warp_perspective.egg-info/PKG-INFO
writing dependency_links to warp_perspective.egg-info/dependency_links.txt
writing top-level names to warp_perspective.egg-info/top_level.txt
writing manifest file 'warp_perspective.egg-info/SOURCES.txt'
reading manifest file 'warp_perspective.egg-info/SOURCES.txt'
writing manifest file 'warp_perspective.egg-info/SOURCES.txt'
running build_ext
copying build/lib.linux-x86_64-3.7/warp_perspective.so ->
Creating /root/local/miniconda/lib/python3.7/site-packages/warp-perspective.egg-link (link to .)
Adding warp-perspective 0.0.0 to easy-install.pth file

Installed /warp_perspective
Processing dependencies for warp-perspective==0.0.0
Finished processing dependencies for warp-perspective==0.0.0

這將產生一個名為 warp_perspective.so 的共享函式庫，我們可以像之前一樣將其傳遞給 torch.ops.load_library，以使我們的運算元對 TorchScript 可見。

>>> import torch
>>> torch.ops.load_library("warp_perspective.so")
>>> print(torch.ops.my_ops.warp_perspective)
<built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>