昨天是介紹如何用 python c/c++ extension 的方法來擴充 PyTorch。今天，我們又回到 TorchScript ，希望能用 TorchScript 來建立一個 C++ operation，然後用 TorchScript 編譯。

關於 TorchScript 我們知道些什麼呢？

1. 在Tensorflow 在 2.0 很好，那麼 PyTorch 呢？這篇文章中，我們知道 TorchScript 在 python 中使用時，有兩種不同的模式，一種是建立靜態計算圖的 tracing module，而另一個則是將流程控制納入編譯好的 Module，稱之為 scripting module。
2. 在TorchScript 實例操作有使用它來 export 在 python 中建置的模型到硬碟內，最後在 C++ 的執行環境中 load 進用 TorchScript 編譯的模組來進行預測。

今天，我們要做的事情幾乎與昨日相等：我們會用 C++ 建立一個新的 operation，昨日為 LLTM，今日則是 warpPerspective 函式，這個函式所做的事情，則是進行視角轉換（perspective transformation）。由於是與電腦影像相關的函式，所以需要 OpenCV 來協助。OpenCV 是用 C++ 寫成關於電腦影像相關的函式庫，它也有 Python binding，不過不在今天的教學內容中使用。

本篇都參考 PyTorch 官方教學：EXTENDING TORCHSCRIPT WITH CUSTOM C++ OPERATORS。若有什麼不明白的，都可以到官方教學網頁仔細閱讀[註一]，也可以在底下留言，大家一起討論。那麼，我們就從 ...

一個具有客製化行為的 C++ 運算元

這個名叫 op.cpp 的檔案，會將 PyTorch 的輸入，也就是具有 torch::Tensor 的型別有下列的流程。

1. 轉換 PyTorch Tensor 物件到 OpenCV 對應的矩陣型態（從 torch::Tensor 到 cv::Mat）。
2. 呼叫 OpenCV 支援函式 warpPerspective。
3. 轉換 OpenCV warpPerspective 輸出為 PyTorch Tensor，如此一來我們可以繼續應用 PyTorch 的函式（從 cv::Mat 到 torch::Tensor）。

步驟 2 是最直覺的，因為就是將 OpenCV 支持的型別丟入 warpPerspective 函式即可。而步驟 1 和 3 則是鏡像的 operations。步驟 1 從 PyTorch Tensor 到 OpenCV Mat，而步驟 2 則從 OpenCV Mat 轉換為 PyTorch Tensor。

為了能使這種轉換更加有效率，盡可能的避免 copy data，在 PyTorch Tensor 端可以使用 torch::from_blob 這個函式。torch::from_blob 需要使用者傳入一個 raw bytes 型態的 buffer（cv::Mat::ptr<float>），這個 buffer 可代表 OpenCV Mat 物件中持有的資料。藉由 OpenCV Mat 物件，release 資料 buffer 的所有權，轉移到 PyTorch Tensor 上，可以避免資料拷貝的問題。

同樣的，關於從 PyTorch Tensor 轉換至 OpenCV Mat 物件，則是利用傳入的 torch::Tensor 物件的 data 屬性來取得物件持有的資料 buffer，交予 OpenCV Mat 的建構子（constructor）即可。

最後一件事，是當函式要回傳資料時，最好使用 clone 的方式，傳回副本。這麼做是因為原函式內若有任何 OpenCV Mat 物件持有該資料 buffer 的參考，資料 buffer 傳回後在原函式外 deallocate 後，會造成 dangling reference 問題。但更尋常的問題，可能是試圖參考一個早在原函式結束時被 deallocate 的空 buffer。

而完整的 C++ 原始碼如下：

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

用 TorchScript 註冊 C++ 運算元

TorchScript 有一個全域的物件， registry，若想透過 TorchScript 去擴充 Torch，就必須以一個特殊的名字註冊入口函式。這個特殊名字通常具有兩個部分，一個是命名空間，另一個則是運算元的函式名稱。這兩個部分，如同其他的命名空間，需要用::隔開。

static auto registry =
  torch::RegisterOperators("my_ops::warp_perspective", &warp_perspective);
  // register more than one operators
  //.op("my_ops::another_op", &another_op) 
  //.op("my_ops::and_another_op", &and_another_op);

torch::RegisterOperators 除了接受一個字串代表的是註冊運算元所在的命名空間，還需要使用者給予實踐該運算元的入口函式指標。同時，若想註冊多個運算元，可以呼叫 .op來完成，如上面的程式碼片段。

Compile and Run

CMake 檔案和專案的檔案結構如下：

warp-perspective/
  op.cpp
  CMakeLists.txt
  build/ # where you invoke cmake command as cmake ../

而 cmake 檔和前幾次的檔案差不多，只是我們需要增加 OpenCV 依賴的部分，包括了標頭檔和函式庫。因為我們的目的是將函式編譯成一個 shared library，所以使用 add_library，而非 add_executable（若用add_executable應該會出錯，因為沒有 main 入口）。檔案內容如下：

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(warp_perspective)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)

# 定義要編譯的函式庫目標
add_library(warp_perspective SHARED op.cpp)
# 使用 C++11
target_compile_features(warp_perspective PRIVATE cxx_range_for)
# 連結到 LibTorch 函式庫
target_link_libraries(warp_perspective "${TORCH_LIBRARIES}")
# 連結到 OpenCV 函式庫
target_link_libraries(warp_perspective opencv_core opencv_imgproc)

若無法編譯成功，則可以使用 PyTorch 官方提供的 Dockerfile。這個 Dockerfile 會使用 apt 安裝 OpenCV。
若成功完成編譯，則會在輸出檔案出現一個 libwarp_perspective.so[註二]，這就是目標 shared library，可以在 python 中使用以下的原始碼，將它 load 進 python 直譯器中：

>>> import torch
>>> torch.ops.load_library("/path/to/libwarp_perspective.so")
>>> print(torch.ops.my_ops.warp_perspective)
#=> <built-in method warp_perspective of PyCapsule object at 0x?????????>
>>> torch.ops.my_ops.warp_perspective(torch.randn(8, 8), torch.rand(3, 3))
#=> tensor([[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  #0.0000],
#        [ 0.0000, -0.0140,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  #0.0000],
#        [ 0.0000, -0.0729, -0.0100,  0.0000,  0.0000,  0.0000,  0.0000,  #0.0000],
#        [ 0.0221, -0.0620, -0.0628,  0.0000,  0.0000,  0.0000,  0.0000,  #0.0000],
#        [-0.0107, -0.0583, -0.0618, -0.0316,  0.0000,  0.0000,  0.0000,  #0.0000],
#        [-0.0111, -0.0583, -0.0610, -0.0465, -0.0057,  0.0000,  0.0000,  #0.0000],
#        [-0.0112, -0.0547, -0.0647, -0.0501, -0.0205,  0.0000,  0.0000,  #0.0000],
#        [-0.0227, -0.0547, -0.0642, -0.0541, -0.0372, -0.0081,  0.0000,  #0.0000]])

成功載入到 python 的直譯器後，我們可以利用 TorchScript 的 JIT 編譯器，這包括了 trace 和 script module。這兩個 modules 會產生相對應的計算圖結構，將程式碼轉變為計算圖。

一旦，我們將完成客製化並編譯成 shared library 的客製化 C++ 運算元，那麼我們就可以在 python 直譯器的 session 內，與其他的 python 原始碼結合，再由 trace 或 script module 編譯成圖。

而關於如何將 TorchScript 編譯結果 export 到磁碟上，在 load 進 C++ 的執行環境，可以參考TorchScript 實例操作。

註釋

[1] 經過測試，需要 PyTorch 1.2 以上的版本，編譯才沒有問題。
[2] 在 MacOSX 環境下編譯成功的共享函示庫檔名為 libwarp_perspective.dylib