使用自定義 C++ 運算元擴充套件 TorchScript¶
建立日期：2018 年 11 月 28 日 | 最後更新：2024 年 7 月 22 日 | 最後驗證：2024 年 11 月 5 日
警告
自 PyTorch 2.4 版本起，本教程已被棄用。請參閱 PyTorch 自定義運算元 獲取最新的 PyTorch 自定義運算元指南。
PyTorch 1.0 版本引入了一種新的程式設計模型，稱為 TorchScript。TorchScript 是 Python 程式語言的一個子集，可以由 TorchScript 編譯器進行解析、編譯和最佳化。此外，編譯後的 TorchScript 模型可以選擇序列化為磁碟檔案格式，隨後你可以從純 C++（以及 Python）載入並執行它進行推理。
TorchScript 支援 torch 包提供的絕大部分操作，這使得你能夠純粹地將許多複雜的模型表示為 PyTorch“標準庫”中的一系列張量操作。然而，有時你可能會發現需要使用自定義 C++ 或 CUDA 函式來擴充套件 TorchScript。雖然我們建議你只在無法（或效率不高）用簡單的 Python 函式表達你的想法時才採用此選項，但我們確實提供了一個非常友好和簡單的介面，用於使用 ATen（PyTorch 的高效能 C++ 張量庫）定義自定義 C++ 和 CUDA 核心。一旦繫結到 TorchScript 中，你可以將這些自定義核心（或“運算元”）嵌入到你的 TorchScript 模型中，並在 Python 中以及以其序列化形式直接在 C++ 中執行它們。
以下段落將舉例說明如何編寫 TorchScript 自定義運算元來呼叫用 C++ 編寫的計算機視覺庫 OpenCV。我們將討論如何在 C++ 中處理張量，如何有效地將它們轉換為第三方張量格式（在本例中是 OpenCV Mat），如何向 TorchScript 執行時註冊你的運算元，最後是如何編譯運算元並在 Python 和 C++ 中使用它。
在 C++ 中實現自定義運算元¶
在本教程中，我們將 OpenCV 的 warpPerspective 函式（該函式對影像應用透視變換）作為自定義運算元暴露給 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。
注意
TorchScript 編譯器僅理解固定數量的型別。只有這些型別才能用作自定義運算元的引數。目前支援的型別有：torch::Tensor、torch::Scalar、double、int64_t 以及它們的 std::vector 型別。請注意，只支援 double 而非 float，只支援 int64_t 而非其他整數型別如 int、short 或 long。
在我們的函式內部，第一件需要做的事情是將 PyTorch 張量轉換為 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 類的 此建構函式 將我們的張量轉換為 Mat 物件。我們向其傳遞了原始 image 張量的行數和列數、資料型別（本例中固定為 float32），最後是一個指向底層資料的原始指標——一個 float*。這個 Mat 建構函式的特殊之處在於它不會複製輸入資料。相反，它只會在所有對 image_mat 執行的操作中引用這塊記憶體。如果在 image_mat 上執行了原地操作，這也會反映在原始的 image 張量中（反之亦然）。這使我們能夠使用庫的原生矩陣型別呼叫後續的 OpenCV 例程，即使我們實際上將資料儲存在 PyTorch 張量中。我們重複此過程將 warp PyTorch 張量轉換為 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 張量，以便我們可以進一步在 PyTorch 中使用它。這與我們之前朝另一個方向轉換時所做的事情非常相似。在這種情況下，PyTorch 提供了一個 torch::from_blob 方法。這裡的 blob 指的是我們想要解釋為 PyTorch 張量的一些不透明的、扁平的記憶體指標。torch::from_blob 的呼叫看起來像這樣：
torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});
return output.clone();
我們使用 OpenCV Mat 類的 .ptr<float>() 方法來獲取指向底層資料的原始指標（就像之前 PyTorch 張量的 .data_ptr<float>()）。我們還指定了張量的輸出形狀，在本例中我們將其硬編碼為 8 x 8。然後，torch::from_blob 的輸出就是一個 torch::Tensor，它指向由 OpenCV 矩陣擁有的記憶體。
在從我們的運算元實現返回這個張量之前，我們必須對該張量呼叫 .clone() 以執行底層記憶體資料的複製。原因是 torch::from_blob 返回的張量並不擁有其資料。那時，資料仍由 OpenCV 矩陣擁有。然而，這個 OpenCV 矩陣在函式結束時會超出作用域並被釋放。如果我們返回的 output 張量原樣不變，那麼在我們函式外部使用它時，它將指向無效的記憶體。呼叫 .clone() 會返回一個新的張量，該張量包含原始資料的副本，並且新張量自己擁有這些資料。因此，將它安全地返回給外部世界是可行的。
向 TorchScript 註冊自定義運算元¶
現在我們已經在 C++ 中實現了自定義運算元，我們需要將其**註冊**到 TorchScript 執行時和編譯器中。這將允許 TorchScript 編譯器解析 TorchScript 程式碼中對我們自定義運算元的引用。如果你使用過 pybind11 庫，我們的註冊語法與其非常相似。要註冊單個函式，我們在 op.cpp 檔案的頂層某處寫入：
TORCH_LIBRARY(my_ops, m) {
m.def("warp_perspective", warp_perspective);
}
TORCH_LIBRARY 宏建立一個函式，該函式會在程式啟動時被呼叫。你的庫的名稱（my_ops）作為第一個引數給出（不應加引號）。第二個引數（m）定義了一個 torch::Library 型別的變數，它是註冊運算元的主要介面。方法 Library::def 實際上建立了一個名為 warp_perspective 的運算元，並將其暴露給 Python 和 TorchScript。你可以透過多次呼叫 def 來定義任意數量的運算元。
在幕後，def 函式實際上做了不少工作：它使用模板超程式設計來檢查函式的型別簽名，並將其轉換為一個運算元 schema，該 schema 在 TorchScript 的型別系統中指定了運算元的型別。
構建自定義運算元¶
現在我們已經在 C++ 中實現了自定義運算元並編寫了註冊程式碼，是時候將運算元構建為一個（共享）庫，以便我們可以在 Python 中載入進行研究和實驗，或者在沒有 Python 的環境中載入到 C++ 中進行推理。存在多種構建運算元する方法，可以使用純 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 使用。雖然此函式可以在 scripted 或 traced TorchScript 模組中使用，但我們也可以在普通的 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 運算元。
從現在開始，你可以在 scripted 或 traced 程式碼中使用你的自定義運算元，就像使用 torch 包中的其他函式一樣。事實上，像 torch.matmul 這樣的“標準庫”函式很大程度上遵循與自定義運算元相同的註冊路徑，這使得自定義運算元在 TorchScript 中的使用方式和位置方面真正成為一等公民。（然而，一個區別是標準庫函式有自定義編寫的 Python 引數解析邏輯，這與 torch.ops 引數解析不同。）
使用 Tracing 自定義運算元¶
讓我們從將我們的運算元嵌入 traced 函式開始。回想一下，對於 tracing，我們從一些普通 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 trace 中，就像它是 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)
然後像之前一樣進行 tracing：
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 自定義運算元整合到 traced PyTorch 程式碼中就是如此簡單！
使用 Script 自定義運算元¶
除了 tracing 之外，獲得 PyTorch 程式 TorchScript 表示的另一種方法是直接**在** TorchScript 中編寫程式碼。TorchScript 很大程度上是 Python 語言的一個子集，帶有一些限制，這些限制使 TorchScript 編譯器更容易理解程式。你透過為自由函式新增 @torch.jit.script 註解，以及為類中的方法新增 @torch.jit.script_method 註解（這些類也必須派生自 torch.jit.ScriptModule）來將你的普通 PyTorch 程式碼轉換為 TorchScript。有關 TorchScript 註解的更多詳細資訊，請參閱 此處。
使用 TorchScript 而非 tracing 的一個特別原因是 tracing 無法捕獲 PyTorch 程式碼中的控制流。因此，讓我們考慮這個使用了控制流的函式：
def compute(x, y):
if bool(x[0][0] == 42):
z = 5
else:
z = 10
return x.matmul(y) + z
要將此函式從普通 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);
}
現在，就像之前一樣，我們可以在 script 程式碼中使用自定義運算元，就像使用任何其他函式一樣：
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 匯入包含你的運算元（或多個運算元）的庫，然後在你的 traced 或 scripted TorchScript 程式碼中像呼叫任何其他 torch 運算元一樣呼叫你的自定義運算元。
在 C++ 中使用 TorchScript 自定義運算元¶
TorchScript 的一個有用特性是能夠將模型序列化到磁碟檔案中。這個檔案可以透過網路傳輸、儲存在檔案系統中，更重要的是，無需保留原始原始碼即可動態反序列化和執行。這在 Python 中是可能的，在 C++ 中也是如此。為此，PyTorch 提供了 一個純 C++ API 用於反序列化和執行 TorchScript 模型。如果你還沒有閱讀過 關於在 C++ 中載入和執行序列化 TorchScript 模型的教程，請先閱讀，接下來的幾段將基於此教程。
簡而言之，自定義運算元可以像常規的 torch 運算元一樣執行，即使是從檔案反序列化並在 C++ 中執行。唯一的必要條件是將我們之前構建的自定義運算元共享庫連結到執行模型的 C++ 應用程式中。在 Python 中，只需呼叫 torch.ops.load_library 即可。在 C++ 中，你需要在使用任何構建系統將共享庫與你的主應用程式連結起來。下面的例子將使用 CMake 來展示這一點。
注意
技術上講，你也可以在 C++ 應用程式執行時動態載入共享庫，方式與我們在 Python 中非常相似。在 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 共享庫連結起來。
注意
上面的示例中隱藏了一個關鍵細節：在 warp_perspective 連結行前面加上了 -Wl,--no-as-needed 字首。這是必需的，因為我們的應用程式程式碼實際上不會呼叫 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++ 運算元擴充套件你的 TorchScript 模型，這些運算元可以與第三方 C++ 庫互動，編寫自定義的高效能 CUDA 核心，或實現任何其他需要 Python、TorchScript 和 C++ 平滑融合的使用場景。
一如既往，如果你遇到任何問題或有疑問，可以使用我們的論壇或GitHub Issues與我們聯絡。此外，我們的常見問題 (FAQ) 頁面可能包含有用的資訊。
附錄 A：更多構建自定義運算元的方法¶
“構建自定義運算元”一節解釋瞭如何使用 CMake 將自定義運算元構建成共享庫。本附錄概述了另外兩種編譯方法。這兩種方法都使用 Python 作為編譯過程的“驅動器”或“介面”。此外，這兩種方法都重用了 PyTorch 為現有基礎設施提供的*C++ 擴充套件*。C++ 擴充套件是 TorchScript 自定義運算元的香草版（eager 模式）對應物，它們依賴於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 中使用此功能，則不應多次執行包含註冊程式碼的單元格，因為每次執行都會註冊一個新的庫並重新註冊自定義運算元。如果需要重新執行，請事先重啟 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>
