用自定義c 運算子擴充套件TorchScript

extending torchscript with custom c++ operators的要點

pytorch 1.0的發布，為 pytorch 引入了乙個新的程式設計模型，稱為 torchscript。torchscript 是 python 的乙個子集，它可以由torchscript 的編譯器進行解析、編譯和優化。進一步地，編譯後的torchscript 模型可以序列化為磁碟上的檔案格式，該檔案格式的推理，可以用純 c + + (以及 python)載入並執行。

torchscript 支援torch包提供的大量運算操作，可以輕易將許多複雜模型表示為pytorch 標準庫中的一系列張量操作。不過，可能會需要擴充套件torchscript 來自定義的 c + + 或 cuda 函式。

下面給出了乙個用 c + +編寫torchscript 自定義程式，實現opencv 的功能

下面實現 warpperspective 函式（對影象進行透視變換），從 opencv 轉化到 torchscript 自定義運算子

第一步是在 c + + 中編寫自定義運算子：

op.cpp

#include #include torch::tensor warp_perspective(torch::tensor image, torch::tensor warp) );
// end output_mat
// begin output_tensor
torch::tensor output = torch::from_blob(output_mat.ptr(), /*sizes=*/);
return output.clone();
// end output_tensor
}

在檔案的頂部，把 opencv 的標頭檔案 opencv2/opencv.hpp和 torch/script.h 都include進來，script.h 裡面公開了 pytorch 的 c + + api 中的所有必需的東西

warp 透視函式接受兩個引數: 乙個是輸入image，另乙個是應用於影象變換的矩陣warp。輸入、返回的型別是 torch::tensor,

注意：torchscript 編譯器理解 torch: : : tensor，torch: : : scalar，double，int64t 和 std: : vector 這些型別。浮點數只支援 double不支援 float，整型只支援 int64_t，不支援 int、 short 或 long 等。

講解：首先將 pytorch 張量轉換為 opencv 矩陣，因為 opencv 的 warpperspective 期望 cv: : mat 作為輸入。

為了不需要複製任何資料:

cv::mat warp_mat(/*rows=*/warp.size(0),
/*cols=*/warp.size(1),
/*type=*/cv_32fc1,
/*data=*/warp.data_ptr());

接下來，在torchscript 中呼叫 opencv的warpperspective。為此，傳入image mat 和 warp mat ，和乙個空的output mat。另外指定了希望輸出矩陣(影象)的dsize為8 x 8:

cv::mat output_mat;
cv::warpperspective(image_mat, output_mat, warp_mat, /*dsize=*/);

自定義運算子實現的最後一步是將output _ mat 轉換到pytorch張量，這樣就可以在 pytorch 中進一步使用它。這與我們之前在另乙個l例子上所做的轉換一致：pytorch提供了乙個 torch: :from_blob 方法，torch: :from_blob的呼叫如下:

torch::tensor output = torch::from_blob(output_mat.ptr(), /*sizes=*/);
return output.clone();

在 opencv 的mat 類上使用 ptr()方法，以獲得乙個指向底層資料的原始指標。from_blob所輸出的 torch::tensor是乙個 torch: : tensor，指向 opencv 矩陣所擁有的記憶體。

對張量進行clone()以執行基礎資料的記憶體副本。原因是 torch: from_blob 返回了乙個沒有其資料的張量，該資料仍然屬於 opencv mat。cv::mat output_mat是區域性變數，在函式結束時釋放，所以呼叫clone()返回乙個新張量，其中包含新張量所擁有的原始資料的副本。

現在已經在 c + + 中實現了自定義運算子，然後需要註冊它。註冊語法與 pybind11語法非常相似：

torch_library(my_ops, m)

在 op.cpp 一開始的某個地方。torch_library 巨集建立乙個在程式啟動時呼叫的函式。庫的名稱(my_ops)作為第乙個引數(不需要引號)。第二個引數(m)定義了乙個型別為 torch: : library 的變數，是註冊運算子的主介面。def 方法實際上建立了乙個名為 warp _ perspective 的操作符，將其暴露給 python 和 torchscript。通過對 def 進行多次呼叫，可以定義任意多個運算子。

有多種方法可以build，這裡使用 setuptools。完全從 python 編譯自定義運算子。優點：setuptools 有乙個非常強大和廣泛的介面，可以用來構建用 c + + 編寫的 python 模組。但是，setuptools 實際上是用於編譯 python的模組，只需要乙個 setup.py 檔案:

from setuptools import setup
from torch.utils.cpp_extension import buildextension, cppextension
setup(
name=
"warp_perspective"
,    ext_modules=
[        cppextension(
"warp_perspective",[
],libraries=
["opencv_core"
,"opencv_imgproc"],
)],    cmdclass=
,)

注意，末尾 buildextension 中啟用了 no_python_abi_suffix ，表示 setuptools 在生成的共享庫名稱中省略任何特定於 python-3的 abi 字尾。例如，在 python 3.7中，這個庫可能被命名為 warp_perspective.cpython-37m-x86_64-linux-gnu。其中cpython-37m-x86_64-linux-gnu 是 abi 標記，但我們實際上只希望它被命名為 warp_perspective。所以如果我們現在在 setup.py 所在的資料夾中執行

python setup.py build develop

得到

running build running build_ext building 'warp_perspective' extension creating build creating build/temp.linux-x86_64-3.7 ...installed /warp_perspective processing dependencies for warp-perspective==0.0.0

finished processing dependencies for warp-perspective==0.0.0

生成乙個叫warp_perspective 的共享庫。然後將其傳遞到 torch.ops.load_library，使運算子對於 torchscript 可見:

torch.ops.load_library(
"warp_perspective.so"
)print
(torch.ops.custom.warp_perspective)
#

用自定義c 運算子擴充套件TorchScript

6 1 2 自定義運算子

6 1 2 自定義運算子

scala自定義運算子

用自定義c 運算子擴充套件TorchScript

6 1 2 自定義運算子

6 1 2 自定義運算子

scala自定義運算子

相關推薦