原文： PyTorch 使用自定義 C ++運算符擴展 TorchScript

PyTorch 1.0 版本向 PyTorch 引入了一種新的編程模型，稱為 TorchScript 。 TorchScript 是 Python 編程語言的子集，可以通過 TorchScript 編譯器進行解析，編譯和優(yōu)化。 此外，已編譯的 TorchScript 模型可以選擇序列化為磁盤文件格式，然后可以從純 C ++(以及 Python）加載并運行該文件格式以進行推理。

TorchScript 支持torch包提供的大量操作子集，使您可以純粹表示為 PyTorch 的“標準庫”中的一系列張量操作來表示多種復雜模型。 但是，有時您可能需要使用自定義 C ++或 CUDA 函數(shù)擴展 TorchScript。 雖然我們建議您僅在無法(簡單有效地）將您的想法表達為簡單的 Python 函數(shù)時才訴諸該選項，但我們確實提供了一個非常友好且簡單的界面，用于使用 ATen 定義自定義 C ++和 CUDA 內核。 ，PyTorch 的高性能 C ++張量庫。 綁定到 TorchScript 后，您可以將這些自定義內核(或“ ops”）嵌入到 TorchScript 模型中，并以 Python 或直接以 C ++的序列化形式執(zhí)行它們。

以下段落提供了編寫 TorchScript 自定義操作以調用 OpenCV (使用 C ++編寫的計算機視覺庫）的示例。 我們將討論如何在 C ++中使用張量，如何有效地將它們轉換為第三方張量格式(在這種情況下為 OpenCV 或 Mat），如何在 TorchScript 運行時中注冊您的運算符 最后是如何編譯運算符并在 Python 和 C ++中使用它。

在 C ++中實現(xiàn)自定義運算符

在本教程中，我們將公開 warpPerspective 函數(shù)，該函數(shù)將透視轉換應用于圖像，從 OpenCV 到 TorchScript 作為自定義運算符。 第一步是用 C ++編寫自定義運算符的實現(xiàn)。 讓我們將此實現(xiàn)的文件稱為op.cpp，并使其如下所示：

#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=*/{8, 8});


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

該運算符的代碼很短。 在文件頂部，我們包含 OpenCV 標頭文件opencv2/opencv.hpp和torch/script.h標頭，該標頭暴露了 PyTorch C ++ API 中所有需要編寫自定義 TorchScript 運算符的必需屬性。 我們的函數(shù)warp_perspective具有兩個參數(shù)：輸入image和我們希望應用于圖像的warp變換矩陣。 這些輸入的類型是torch::Tensor，這是 C ++中 PyTorch 的張量類型(也是 Python 中所有張量的基礎類型）。 我們的warp_perspective函數(shù)的返回類型也將是torch::Tensor。

小費

有關 ATen 的更多信息，請參見本說明，ATen 是為 PyTorch 提供Tensor類的庫。 此外，本教程的描述了如何在 C ++中分配和初始化新的張量對象(此運算符不需要）。

注意

TorchScript 編譯器了解固定數(shù)量的類型。 只有這些類型可以用作自定義運算符的參數(shù)。 當前這些類型是：這些類型的torch::Tensor，torch::Scalar，double，int64_t和std::vector。 請注意，僅，，double和不，，float，僅，，int64_t和，等其他整數(shù)類型，例如int 支持short或long。

在函數(shù)內部，我們要做的第一件事是將 PyTorch 張量轉換為 OpenCV 矩陣，因為 OpenCV 的warpPerspective期望cv::Mat對象作為輸入。 幸運的是，有一種方法可以執(zhí)行此，而無需復制任何數(shù)據(jù)。 在前幾行中

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

我們正在將稱為 OpenCV Mat類的構造函數(shù)，將張量轉換為Mat對象。 我們將原始image張量的行數(shù)和列數(shù)，數(shù)據(jù)類型(在此示例中，我們將其固定為float32）傳遞給它，最后傳遞指向基礎數(shù)據(jù)的原始指針– float*。 Mat類的此構造方法的特殊之處在于它不會復制輸入數(shù)據(jù)。 取而代之的是，它將簡單地引用此內存來執(zhí)行Mat上的所有操作。 如果在image_mat上執(zhí)行就地操作，這將反映在原始image張量中(反之亦然）。 即使我們實際上將數(shù)據(jù)存儲在 PyTorch 張量中，這也使我們能夠使用庫的本機矩陣類型調用后續(xù)的 OpenCV 例程。 我們重復此過程將warp PyTorch 張量轉換為warp_mat OpenCV 矩陣：

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

接下來，我們準備調用我們渴望在 TorchScript 中使用的 OpenCV 函數(shù)：warpPerspective。 為此，我們將image_mat和warp_mat矩陣以及稱為output_mat的空輸出矩陣傳遞給 OpenCV 函數(shù)。 我們還指定了我們希望輸出矩陣(圖像）為dsize的大小。 對于此示例，它被硬編碼為8 x 8：

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

我們的自定義運算符實現(xiàn)的最后一步是將output_mat轉換回 PyTorch 張量，以便我們可以在 PyTorch 中進一步使用它。 這與我們先前在另一個方向進行轉換的操作極為相似。 在這種情況下，PyTorch 提供了torch::from_blob方法。 在這種情況下， blob 旨在表示一些不透明的，扁平的指向內存的指針，我們希望將其解釋為 PyTorch 張量。 對torch::from_blob的調用如下所示：

torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8})

我們在 OpenCV Mat類上使用.ptr<float>()方法來獲取指向基礎數(shù)據(jù)的原始指針(就像之前的 PyTorch 張量的.data<float>()一樣）。 我們還指定了張量的輸出形狀，我們將其硬編碼為8 x 8。 然后torch::from_blob的輸出是torch::Tensor，指向 OpenCV 矩陣擁有的內存。

從我們的運算符實現(xiàn)返回該張量之前，我們必須在張量上調用.clone()以執(zhí)行基礎數(shù)據(jù)的存儲副本。 這樣做的原因是torch::from_blob返回的張量不擁有其數(shù)據(jù)。 那時，數(shù)據(jù)仍歸 OpenCV 矩陣所有。 但是，此 OpenCV 矩陣將超出范圍，并在函數(shù)末尾重新分配。 如果我們按原樣返回output張量，那么當我們在函數(shù)外使用它時，它將指向無效的內存。 調用.clone()將返回一個新的張量，其中包含新張量自己擁有的原始數(shù)據(jù)的副本。 因此，返回外部世界是安全的。

使用 TorchScript 注冊自定義運算符

現(xiàn)在，已經在 C ++中實現(xiàn)了自定義運算符，我們需要在 TorchScript 運行時和編譯器中將注冊為。 這將使 TorchScript 編譯器可以在 TorchScript 代碼中解析對我們自定義運算符的引用。 注冊非常簡單。 對于我們的情況，我們需要編寫：

static auto registry =
  torch::RegisterOperators("my_ops::warp_perspective", &warp_perspective);

op.cpp文件的全局范圍內的某個位置。 這將創(chuàng)建一個全局變量registry，該變量將在其構造函數(shù)中向 TorchScript 注冊我們的運算符(即每個程序一次）。 我們指定運算符的名稱，以及指向其實現(xiàn)的指針(我們之前編寫的函數(shù)）。 該名稱包括兩部分：命名空間(my_ops）和我們正在注冊的特定運算符的名稱(warp_perspective）。 名稱空間和操作員名稱由兩個冒號(::）分隔。

Tip

如果要注冊多個運算符，可以在構造函數(shù)之后將調用鏈接到.op()：

static auto registry =
  torch::RegisterOperators("my_ops::warp_perspective", &warp_perspective)
  .op("my_ops::another_op", &another_op)
  .op("my_ops::and_another_op", &and_another_op);

在后臺，RegisterOperators將執(zhí)行許多相當復雜的 C ++模板元編程魔術技巧，以推斷我們傳遞給它的函數(shù)指針的參數(shù)和返回值類型(&warp_perspective）。 此信息用于為我們的操作員形成功能模式。 函數(shù)模式是操作員的結構化表示形式，一種“簽名”或“原型”，由 TorchScript 編譯器用來驗證 TorchScript 程序的正確性。

建立自定義操作員

現(xiàn)在，我們已經用 C ++實現(xiàn)了自定義運算符并編寫了其注冊代碼，是時候將該運算符構建到一個(共享的）庫中了，可以將其加載到 Python 中進行研究和實驗，或者加載到 C ++中以在非 Python 中進行推理。 環(huán)境。 有多種方法可以使用純 CMake 或setuptools之類的 Python 替代方法來構建我們的運算符。 為簡潔起見，以下段落僅討論 CMake 方法。 本教程的附錄深入探討了基于 Python 的替代方法。

用 CMake 構建

為了使用 CMake 構建系統(tǒng)將自定義運算符構建到共享庫中，我們需要編寫一個簡短的CMakeLists.txt文件并將其與之前的op.cpp文件一起放置。 為此，讓我們就一個看起來像這樣的目錄結構達成一致：

warp-perspective/
  op.cpp
  CMakeLists.txt

另外，請確保從 pytorch.org 中獲取 LibTorch 發(fā)行版的最新版本，該軟件包打包了 PyTorch 的 C ++庫和 CMake 構建文件。 將解壓縮的發(fā)行版放置在文件系統(tǒng)中可訪問的位置。 以下段落將將該位置稱為/path/to/libtorch。 我們的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++11
target_compile_features(warp_perspective PRIVATE cxx_range_for)
## Link against LibTorch
target_link_libraries(warp_perspective "${TORCH_LIBRARIES}")
## Link against OpenCV
target_link_libraries(warp_perspective opencv_core opencv_imgproc)

警告

此設置對構建環(huán)境進行了一些假設，特別是有關 OpenCV 安裝的假設。 上面的CMakeLists.txt文件已在運行 Ubuntu Xenial 的 Docker 容器中通過apt安裝了libopencv-dev進行了測試。 如果它對您不起作用，并且您感到困惑，請使用隨附的教程資料庫中的Dockerfile構建一個隔離的，可復制的環(huán)境，在其中可以使用本教程中的代碼。 如果您遇到其他麻煩，請在教程資料庫中提交問題，或在我們的論壇中發(fā)布問題。

現(xiàn)在要構建我們的操作員，我們可以從warp_perspective文件夾中運行以下命令：

$ mkdir build
$ cd 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
-- 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命令中，應將/path/to/libtorch替換為未壓縮的 LibTorch 發(fā)行版的路徑。

我們將在下面進一步探討如何使用和調用我們的運算符，但是為了早日獲得成功，我們可以嘗試在 Python 中運行以下代碼：

>>> import torch
>>> torch.ops.load_library("/path/to/libwarp_perspective.so")
>>> print(torch.ops.my_ops.warp_perspective)

在這里，/path/to/libwarp_perspective.so應該是我們剛剛構建的libwarp_perspective.so共享庫的相對或絕對路徑。 如果一切順利，這應該打印類似

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

這是我們稍后將用來調用自定義運算符的 Python 函數(shù)。

在 Python 中使用 TorchScript 自定義運算符

將我們的自定義運算符構建到共享庫后，我們就可以在 Python 的 TorchScript 模型中使用此運算符了。 這有兩個部分：首先將運算符加載到 Python 中，其次在 TorchScript 代碼中使用運算符。

您已經了解了如何將運算符導入 Python：torch.ops.load_library()。 此函數(shù)采用包含自定義運算符的共享庫的路徑，并將其加載到當前進程中。 加載共享庫還將執(zhí)行我們放入自定義運算符實現(xiàn)文件中的全局RegisterOperators對象的構造函數(shù)。 這將在 TorchScript 編譯器中注冊我們的自定義運算符，并允許我們在 TorchScript 代碼中使用該運算符。

您可以將已加載的運算符稱為torch.ops.<namespace>.<function>，其中<namespace>是運算符名稱的名稱空間部分，而<function>是運算符的函數(shù)名稱。 對于我們上面編寫的運算符，名稱空間為my_ops，函數(shù)名稱為warp_perspective，這意味著我們的運算符可以作為torch.ops.my_ops.warp_perspective使用。 盡管可以在腳本化或跟蹤的 TorchScript 模塊中使用此函數(shù)，但我們也可以僅在原始的 PyTorch 中使用它，并將其傳遞給常規(guī) PyTorch 張量：

>>> import torch
>>> torch.ops.load_library("libwarp_perspective.so")
>>> 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]])

注意

幕后發(fā)生的事情是，第一次使用 Python 訪問torch.ops.namespace.function時，TorchScript 編譯器(在 C ++平臺上）將查看是否已注冊函數(shù)namespace::function，如果已注冊，則將 Python 句柄返回給該函數(shù)， 我們可以隨后使用它從 Python 調用我們的 C ++運算符實現(xiàn)。 這是 TorchScript 自定義運算符和 C ++擴展之間的一個值得注意的區(qū)別：C ++擴展是使用 pybind11 手動綁定的，而 TorchScript 自定義操作則是由 PyTorch 自己動態(tài)綁定的。 Pybind11 在綁定到 Python 的類型和類方面為您提供了更大的靈活性，因此建議將其用于純粹渴望的代碼，但 TorchScript ops 不支持它。

從這里開始，您可以在腳本或跟蹤代碼中使用自定義運算符，就像torch包中的其他函數(shù)一樣。 實際上，諸如torch.matmul之類的“標準庫”功能與自定義運算符的注冊路徑大致相同，這使得自定義運算符在 TorchScript 中的使用方式和位置方面真正成為一等公民。

使用自定義運算符進行跟蹤

首先，將我們的運算符嵌入到跟蹤函數(shù)中。 回想一下，為了進行跟蹤，我們從一些原始的 Pytorch 代碼開始：

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

然后調用torch.jit.trace。 我們進一步傳遞torch.jit.trace一些示例輸入，它將輸入到我們的實現(xiàn)中，以記錄輸入流過它時發(fā)生的操作順序。 這樣的結果實際上是渴望的 PyTorch 程序的“凍結”版本，TorchScript 編譯器可以對其進行進一步的分析，優(yōu)化和序列化：

>>> 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)
    %y : Float(8, 5)
    %z : Float(4, 5)) {
  %3 : Float(4, 5) = aten::matmul(%x, %y)
  %4 : Float(4, 5) = aten::relu(%z)
  %5 : int = prim::Constant[value=1]()
  %6 : Float(4, 5) = aten::add(%3, %4, %5)
  return (%6);
}

現(xiàn)在，令人興奮的啟示是，我們可以簡單地將自定義運算符放到 PyTorch 跟蹤中，就好像它是torch.relu或任何其他torch函數(shù)一樣：

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


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)
    %y : Float(8, 5)
    %z : Float(8, 5)) {
    %3 : int = prim::Constant[value=3]()
    %4 : int = prim::Constant[value=6]()
    %5 : int = prim::Constant[value=0]()
    %6 : int[] = prim::Constant[value=[0, -1]]()
    %7 : Float(3, 3) = aten::eye(%3, %4, %5, %6)
    %x : Float(8, 8) = my_ops::warp_perspective(%x.1, %7)
    %11 : Float(8, 5) = aten::matmul(%x, %y)
    %12 : Float(8, 5) = aten::relu(%z)
    %13 : int = prim::Constant[value=1]()
    %14 : Float(8, 5) = aten::add(%11, %12, %13)
    return (%14);
  }

如此簡單地將 TorchScript 自定義操作集成到跟蹤的 PyTorch 代碼中！

將自定義運算符與腳本一起使用

除了跟蹤之外，獲得 PyTorch 程序的 TorchScript 表示形式的另一種方法是直接在 TorchScript 中編寫代碼。 TorchScript 在很大程度上是 Python 語言的子集，它具有一些限制，使 TorchScript 編譯器更容易推理程序。 您可以使用@torch.jit.script標記免費功能，使用@torch.jit.script_method標記類中的方法(也必須從torch.jit.ScriptModule派生），將常規(guī) PyTorch 代碼轉換為 TorchScript。 有關 TorchScript 注釋的更多詳細信息，請參見此處的。

使用 TorchScript 而不是跟蹤的一個特殊原因是，跟蹤無法捕獲 PyTorch 代碼中的控制流。 因此，讓我們考慮使用控制流的此函數(shù)：

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函數(shù)編譯為圖形表示形式，我們可以在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);
}

現(xiàn)在，就像以前一樣，我們可以像腳本代碼中的任何其他函數(shù)一樣使用自定義運算符：

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 ++中的RegisterOperators對象注冊的實現(xiàn)，并將其編譯為圖形表示形式：

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

請?zhí)貏e注意圖形末尾對my_ops::warp_perspective的引用。

Attention

TorchScript 圖形表示仍可能更改。 不要依靠它看起來像這樣。

在 Python 中使用自定義運算符時，確實如此。 簡而言之，您可以使用torch.ops.load_library導入包含運算符的庫，并像其他任何torch運算符一樣，從跟蹤或編寫腳本的 TorchScript 代碼中調用自定義操作。

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

TorchScript 的一項有用功能是能夠將模型序列化到磁盤文件中。 該文件可以通過有線方式發(fā)送，存儲在文件系統(tǒng)中，或者更重要的是，可以動態(tài)反序列化和執(zhí)行，而無需保留原始源代碼。 這在 Python 中是可能的，但在 C ++中也是可能的。 為此，PyTorch 為提供了純 C ++ API ，用于反序列化以及執(zhí)行 TorchScript 模型。 如果還沒有的話，請閱讀有關使用 C ++ 加載和運行序列化 TorchScript 模型的教程，接下來的幾段將基于該教程構建。

簡而言之，即使從文件反序列化并以 C ++運行，也可以像常規(guī)torch運算符一樣執(zhí)行自定義運算符。 唯一的要求是將我們先前構建的自定義運算符共享庫與執(zhí)行模型的 C ++應用程序鏈接。 在 Python 中，只需調用torch.ops.load_library即可。 在 C ++中，您需要在使用的任何構建系統(tǒng)中將共享庫與主應用程序鏈接。 下面的示例將使用 CMake 展示這一點。

Note

從技術上講，您還可以在運行時將共享庫動態(tài)加載到 C ++應用程序中，就像在 Python 中一樣。 在 Linux 上，可以使用 dlopen 來執(zhí)行此操作。 在其他平臺上也存在等效項。

在上面鏈接的 C ++執(zhí)行教程的基礎上，讓我們從一個文件中的最小 C ++應用程序開始，該文件位于與自定義運算符不同的文件夾中的main.cpp，該文件加載并執(zhí)行序列化的 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().
  std::shared_ptr<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=/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
-- 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>

接下來，讓我們序列化我們先前編寫的使用自定義運算符的腳本函數(shù)：

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)

或者可能不是。 也許還沒有。 當然！ 我們尚未將自定義運算符庫與我們的應用程序鏈接。 讓我們立即執(zhí)行此操作，并正確進行操作，讓我們稍微更新一下文件組織，如下所示：

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共享庫鏈接起來。

Attention

上面的示例中嵌入了一個關鍵細節(jié)：warp_perspective鏈接行的-Wl,--no-as-needed前綴。 這是必需的，因為我們實際上不會在應用程序代碼中從warp_perspective共享庫中調用任何函數(shù)。 我們只需要運行全局RegisterOperators對象的構造函數(shù)即可。 麻煩的是，這使鏈接器感到困惑，并使其認為可以完全跳過針對庫的鏈接。 在 Linux 上，-Wl,--no-as-needed標志強制執(zhí)行鏈接(注意：該標志特定于 Linux?。?還有其他解決方法。 最簡單的方法是在操作員庫中定義一些函數(shù)，您需要從主應用程序中調用該函數(shù)。 這可能就像在某個標頭中聲明的函數(shù)void init();一樣簡單，然后在運算符庫中將其定義為void init() { }。 在主應用程序中調用此init()函數(shù)會給鏈接器以印象，這是一個值得鏈接的庫。 不幸的是，這不在我們的控制范圍之內，我們寧愿讓您知道其原因和簡單的解決方法，而不是讓您將一些不透明的宏放入代碼中。

現(xiàn)在，由于我們現(xiàn)在在頂層找到了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=/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
-- 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

如果現(xiàn)在運行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} ]

成功！ 您現(xiàn)在可以推斷了。

結論

本教程向您介紹了如何在 C ++中實現(xiàn)自定義 TorchScript 運算符，如何將其構建到共享庫中，如何在 Python 中使用它來定義 TorchScript 模型，最后如何將其加載到 C ++應用程序中以進行推理工作負載。 現(xiàn)在，您可以使用與第三方 C ++庫進行接口的 C ++運算符擴展 TorchScript 模型，編寫自定義的高性能 CUDA 內核，或實現(xiàn)任何其他需要 Python，TorchScript 和 C ++之間的界線才能平穩(wěn)融合的用例。

與往常一樣，如果您遇到任何問題或疑問，可以使用我們的論壇或 GitHub 問題進行聯(lián)系。 另外，我們的常見問題解答(FAQ）頁面可能包含有用的信息。

附錄 A：建立自定義操作員的更多方法

“構建自定義運算符”一節(jié)介紹了如何使用 CMake 將自定義運算符構建到共享庫中。 本附錄概述了兩種進一步的編譯方法。 他們倆都使用 Python 作為編譯過程的“驅動程序”或“接口”。 此外，兩者都重新使用了現(xiàn)有基礎結構 PyTorch 提供了 C ++擴展 ，它們是依賴于 [pybind11 用于將功能從 C ++“顯式”綁定到 Python。

第一種方法是使用 C ++擴展程序的方便的即時(JIT）編譯界面在您首次運行 PyTorch 腳本時在后臺編譯代碼。 第二種方法依賴于古老的setuptools包，并涉及編寫單獨的setup.py文件。 這樣可以進行更高級的配置，并與其他基于setuptools的項目集成。 我們將在下面詳細探討這兩種方法。

使用 JIT 編譯進行構建

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

Note

這里的“ JIT 編譯”與 TorchScript 編譯器中用于優(yōu)化程序的 JIT 編譯無關。 這只是意味著您的自定義運算符 C ++代碼將在您首次導入時在系統(tǒng) <cite>/ tmp</cite> 目錄下的文件夾中編譯，就像您自己事先對其進行編譯一樣。

此 JIT 編譯功能有兩種形式。 首先，您仍然將操作員實現(xiàn)保存在單獨的文件(op.cpp）中，然后使用torch.utils.cpp_extension.load()編譯擴展名。 通常，此函數(shù)將返回暴露您的 C ++擴展的 Python 模塊。 但是，由于我們沒有將自定義運算符編譯到其自己的 Python 模塊中，因此我們只想編譯一個普通的共享庫。 幸運的是，torch.utils.cpp_extension.load()有一個參數(shù)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)

這應該大致打?。?/p>

<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();
}


static auto registry =
  torch::RegisterOperators("my_ops::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 中使用此功能，則不應多次執(zhí)行單元格的注冊，因為每次執(zhí)行都會注冊一個新庫并重新注冊自定義運算符。 如果需要重新執(zhí)行它，請事先重新啟動筆記本的 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

如果現(xiàn)在從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.custom.warp_perspective)
<built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>