Python 扩展 Op
Python 擴展 Op
注意 ：本文涉及的 Python Kernel 僅在 gcc 4.8.5 編譯環境下充分測試，進一步的完善計劃見 Issue 3951。
背景介紹
OneFlow 將各種對于數據的處理都抽象成了算子（operator），簡稱 op。 op 是作用在輸入 tensor 上的操作，并將操作的結果寫到輸出 tensor 上。OneFlow 內部已經提供了比較完備的 op 算子，可以在 ops 目錄下找到。
當 OneFlow 已有的 Python 算子及其組合無法滿足構建神經網絡的需求，或者 Python 層次的算子無法滿足性能需求時，可以開發自定義 op。OneFlow 提供了兩類開發自定義 Op 的途徑，一類是以 Python 為主的 Python Kernel 開發，另外一類是使用 C++ 擴展 Op一文介紹的 C++ Kernel 開發。
Python Kernel 因為主要采用 Python 進行擴展，開發流程較簡單，適用于快速預研、算法驗證等場景。C++ Kernel 效率高，適用于開發已經驗證穩健性并追求性能的算子。
本文介紹算子開發的背景知識和基本概念，并展示如何開發 Python Kernel。
基本概念
在進行 OneFlow 算子開發前，需要了解 op_type_name、Op 以及 Kernel 這幾個概念：
? op_type_name：op_type_name 是 op 類別的全局唯一 ID， OneFlow 通過 op_type_name 查詢并確認 op 的種類，進而實例化 op，用于構建計算圖。op 的種類與 op 的關系，類似于類與對象的關系。
? op：邏輯上的算子，包含構圖推理時的輸入輸出形狀等信息，不包含具體的處理數據的邏輯。
? kernel：對于一個邏輯上的 op，在運行時，處理的邏輯會因為物理設備以及數據類型的不同。運行時的具體處理邏輯，由 kernel 完成。簡單而言，op 與 kernel 是一對多的關系，可以使用 Python 完成具體運算，這樣的Kernel 稱為 Python Kernel，也可以使用 C++ 開發 Kernel。
? OneFlow 的內核由 C++ 實現，但是用戶接口使用 Python，需要按照約定編寫 Python Wrapper，使得 Python Op 接口能與 C++ 內核交互。
開發步驟
使用 Python 擴展 Op，應該準備一個以 op_type_name 命名的目錄，在該目錄下，按照約定放置必需的文件，以 oneflow/python/test/custom_ops/user_sigmoid 為例：
user_sigmoid
├── user_sigmoid_cpp_def.cpp
├── user_sigmoid_py_api.py
└── user_sigmoid_py_kernel.py
其中：
? op_type_name_cpp_def.cpp(以上的 user_sigmoid_cpp_def.cpp) 文件中放置 Op 定義信息
? op_type_name_py_api.py(以上的 user_sigmoid_py_api.py)文件中放置 Python Wrapper，通過 oneflow.user_op_builder 將實現的 Python Kernel 導出給用戶使用
? op_type_name_py_kernel.py(以上的 user_sigmoid_py_kernel.py)文件中放置 Python 實現的自定義算子的前向計算邏輯和后向計算邏輯
在下文中，將分別介紹：
? 如何編寫 op_type_name_cpp_def.cpp 文件，定義 Op 信息
? 如何編寫 op_type_name_py_api.py 文件，封裝 Op 的 Python 接口
? 如何編寫 op_type_name_py_kernel.py 文件，使用 Python 實現 Op 的計算 Kernel
? 在 OneFlow 中如何使用 Python Kernel 類型的自定義 Op
下文中，將介紹如何用 Python 實現一個自定義的 user_relu Op。
Op 的實現與注冊
首先，在 user_relu_cpp_def.cpp 中定義 op 并完成注冊：
#include “oneflow/core/framework/framework.h”
namespace oneflow {
namespace {
REGISTER_USER_OP(“user_relu_forward”)
.Attrstd::string(“device_sub_tag”, “py”)
.Input(“in”)
.Output(“out”)
.SetTensorDescInferFn(
[](user_op::InferContext *ctx) -> Maybe {
*ctx->Shape4ArgNameAndIndex(“out”, 0) =
*ctx->Shape4ArgNameAndIndex(“in”, 0);
*ctx->Dtype4ArgNameAndIndex(“out”, 0) =
*ctx->Dtype4ArgNameAndIndex(“in”, 0);
return Maybe::Ok();
});
} // namespace
} // namespace oneflow
分析以上代碼：
? oneflow/core/framework/framework.h 中包含了創建一個 op 所需要的所有接口
? .Attrstd::string(“device_sub_tag”, “py”) 是必需的，它告知 OneFlow 在使用該 Op 時默認調用Python Kernel
? 與自定義 op 有關的接口集中在 oneflow::user_op 中，使用名稱空間 oneflow 可以簡化類型名稱
? 宏 REGISTER_USER_OP 用于注冊 op，其接受的參數 user_relu_forward 是 op_type_name。
? 使用 REGISTER_USER_OP 注冊后，其實會返回一個 OpRegistry 類（位于user_op_registry.h），通過調用該類方法，完成對自定義 op 的設置：
a. Input(“in”) 表示其有一個名為 “in” 的輸入
b. Output(“out”) 表示其有一個名為 “out” 的輸出
c. SetTensorDescInferFn 用于設置形狀及數據類型推導函數，描述該算子的輸出的形狀及類型與輸入的關系。以上代碼中，輸出的形狀、數據類型與輸入的一致
op_type_name_cpp_def.cpp 文件是實現 Python Kernel 過程中唯一會使用到的 C++ 文件，用于設置 Op 的信息，在現階段，還無法將使用 C++ 配置 Op 的步驟省略（因為設置分布式等高級信息時必需），不過可以看到，該文件并不涉及具體的運算，僅僅是用于描述 Op，即使不熟悉 C++，根據的示例，也可以很輕松地掌握。
封裝 Op 的 Python 接口
為了用戶可以在 Python 層使用剛剛設置并注冊的 user_relu Op，需要創建一個 user_relu_py_api.py 文件，其內容如下：
import oneflow as flow
def user_relu_forward(x):
op = (
flow.user_op_builder(“myrelu”)
.Op(“user_relu_forward”)
.Input(“in”, [x])
.Output(“out”)
.Build()
)
return op.InferAndTryRun().SoleOutputBlob()
flow.user_op_builder(“op_myrelu”) 其實會返回一個名為 op_myrelu 的 UserOpConfBuilder 對象。
該對象包含 Op、Input 等方法，用于封裝自定義 op，具體解釋如下：
? Op(“user_relu_forward”)：參數必須為之前在 C++ 注冊時的 op_type_name，OneFlow 通過它找到已經注冊的 op 類型，并實例化 op 對象。
? Input(“in”, [input_blob])：對應了 C++ 中 op 注冊時的 Input，第一個參數字符串必須與 C++ 注冊 op 時的 Input 設置的字符串一致。第二個參數為輸入的張量，是一個 list，因為一個 op 允許有多個輸入。
? Output(“out”)：對應了 C++ 中 op 注冊時的 Output。
? Build：以上設置完成后，調用 Build 可以得到自定義 op 的 Python wrapper
以下代碼，將獲取自定義 op 的輸出：
return op.InferAndTryRun().SoleOutputBlob()
其中的 InferAndTryRun 完成推導，返回 UserOp，如果返回結果只有一個輸出，則使用 SoleOutputBlob 即可獲取該唯一輸出，否則，可以使用 RemoteBlobList 獲取包含多個輸出的列表。
使用 Python 實現 Kernel
如本文開始所描述，Op 只是邏輯上的概念，真正的計算需要 Kernel 完成，在 OneFlow 中可以既可以使用 C++ ，也可以使用 Python 實現 Kernel，本文只介紹最易上手的 Python Kernel 的實現方法。使用 C++ 實現 Kernel 可以參考使用 C++ 開發 Kernel。
為了為上文設置的 user_relu Op 提供 Python Kernel，需要創建一個 user_relu_py_kernel.py 文件，其內容如下：
import numpy as np
def forward(args):
(x,) = args
y = (x>0)*x
return y
以上的 forward 方法是必需實現的，它的實現對應了 Op 的 Python Kernel。關于它的約定有：
? 方法名必需為 forward
? 參數只有一個，類型為 tuple，tuple 中的元素個數和順序，與 Op 注冊時的 Input 對應。如之前為 user_relu 注冊了 Input(“in”)，那么以上代碼中 (x, ) = args 中的 x 就取到 in 的值
? 輸出與 Op 注冊時的 Output 對應
? 參數與返回值均為 numpy 對象，即不能（不會）是字符串、整型數字等其它類型
使用自定義 Op
完成以上工作后，得到了一個名為 user_relu 的目錄，包含三個文件，它們的結構如下：
user_relu/
├── user_relu_cpp_def.cpp
├── user_relu_py_api.py
└── user_relu_py_kernel.py
可以在 user_relu 文件夾所在的路徑，創建一個測試文件，調用剛剛實現的自定義 Op，內容如下：
import oneflow as flow
import numpy as np
import os
import oneflow.typing as tp
根據指定的路徑與 op_type_name 創建 module 對象
module_path = os.path.dirname(os.path.abspath(file))
user_relu_op = flow.experimental.custom_op_module(“user_relu”, module_path)
使 Op, Python API, Python Kernel 生效
user_relu_op.py_api().cpp_def().py_kernel().build_load()
@flow.global_function()
def MyJob(x: tp.Numpy.Placeholder((5,), dtype=flow.float32)) -> tp.Numpy:
with flow.scope.placement(“cpu”, “0:0”):
return user_relu_op.api.user_relu_forward(x)
if name == “main”:
input = np.array([-2, -1, 0, 1, 2], dtype=np.float32)
output = MyJob(input)
print(input)
print(output)
以上代碼中，先通過 flow.experimental.custom_op_module 創建 module 對象，它接收兩個參數，第一個參數為 op_type_name， 第二個參數為 user_relu 文件夾所在的路徑。返回的 module 對象，代表了自定義的 Op。
接著，通過 user_sigmoid_op.py_api().cpp_def().py_kernel().build_load() 可以使自定義 Op 生效，生效后的 Op 的 Python 接口，就是定義在 user_relu_py_api.py 文件中的方法名(user_relu_forward)，它被放置在 moudle 對象的 api 名稱空間中。需要通過以下方式調用:
user_sigmoid_op.api.user_relu_forward(x)
且因為 Python Kernel 只能運行在 CPU 設備上，需要指定計算設備為 CPU：
with flow.scope.placement(“cpu”, “0:0”):
為自定義 Op 提供反向計算
通過上述工作，已經完成了 user_relu 算子的正向計算過程，可以用于 type=“predict” 的作業函數。如果想支持 type=“train” 類型的訓練作業函數，需要為自定義 Op 提供反向計算。
自定義 Op 提供反向計算的代碼，需要寫在 op_type_name_cpp_def.cpp 文件中，通過宏 REGISTER_USER_OP_GRAD 進行注冊。
從數學角度上看，注冊過程就是為自定義的 op，指定后向求梯度的計算方法。從編程角度看，就是為自定義 op 設置一個后向生成函數，在該函數中，編寫代碼，指定這個 op 的輸入梯度的計算方法。
以下，將專門實現一個 Op，名為 user_relu_backward。將在為 user_relu 注冊后向梯度時，用到這個“專門定制”的 Op。
實現 user_relu_backward Op
實現 user_relu_backward Op 的過程與實現 user_relu 的前向幾乎是一樣的。首先，在 user_relu_cpp_def.cpp 中設置并注冊該 Op：
REGISTER_USER_OP(“user_relu_backward”)
.Input(“y”)
.Input(“dy”)
.Output(“dx”)
.Attrstd::string(“device_sub_tag”, “py”)
.SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe {
const Shape* dy_shape = ctx->Shape4ArgNameAndIndex(“dy”, 0);
Shape* dx_shape = ctx->Shape4ArgNameAndIndex(“dx”, 0);
*dx_shape = *dy_shape;
return Maybe::Ok();
});
值得注意的是，同前向類似，以上代碼中 .Attrstd::string(“device_sub_tag”, “py”) 必不可少，它告知 OneFlow 在使用該 Op 時，默認調用 Python Kernel。
同理，不需要用戶直接調用這個 user_relu_backward Op，不需要在 user_relu_py_api.py 為 user_relu_backward 封裝 Python 接口。可以直接實現它的 Python Kernel。
在 user_relu_py_kernel.py 中，實現 backward 方法：
def backward(args):
(y, dy) = args
dx = (y>0)dy
return dx
它的參數是一個 tuple，數目和順序對應了 Op 注冊時的 Input，輸出對應了 Op 注冊時的 Output。
為 Op 注冊反向梯度
需要在 user_relu_cpp_def.cpp 中，通過宏 REGISTER_USER_OP_GRAD 為的正向 Op (user_relu_forward) 注冊反向。
其代碼如下：
REGISTER_USER_OP_GRAD(“user_relu_forward”)
.SetBackwardOpConfGenFn([](user_op::BackwardOpConfContext ctx) {
const auto grad_op_name = ctx->FwOp().op_name() + “_grad”;
const auto& grad_op_func = [&ctx](user_op::BackwardOpBuilder& builder) {
return builder.OpTypeName(“user_relu_backward”)
.InputBind(“y”, ctx->FwOp().output(“y”, 0))
.InputBind(“dy”, ctx->FwOp().output_grad(“y”, 0))
.Output(“dx”)
.Build();
};
ctx->DefineOp(grad_op_name, grad_op_func);
const auto& dx_get_func = [&ctx, &grad_op_name]() -> const std::string& {return ctx->GetOp(grad_op_name).output("dx", 0);};ctx->FwOp().InputGradBind(user_op::OpArg("x", 0), dx_get_func);
});
對以上代碼進行解釋，通過 REGISTER_USER_OP_GRAD(“user_relu_forward”) 注冊為前向 Op 注冊后向求梯度規則，該宏接收一個參數，就是 前向的 op_type_name。
然后通過 SetBackwardOpConfGenFn 設置后向求梯度規則，同 Op 類似，在 op_type_name_cpp_def.cpp 中注冊后向，其實不涉及真正的運算，而是設置后向計算與前向的對應關系，告訴 OneFlow 框架：
? 用什么 Op 求后向梯度
? 該 Op 的輸入來自哪里，和前向 Op 什么關系
因此，以上代碼中的：
const auto& grad_op_func = [&ctx](user_op::BackwardOpBuilder& builder) {
return builder.OpTypeName(“user_relu_backward”)
.InputBind(“y”, ctx->FwOp().output(“y”, 0))
.InputBind(“dy”, ctx->FwOp().output_grad(“y”, 0))
.Output(“dx”)
.Build();
};
定義了 Op 求梯度的方法：使用 user_relu_backward 算子，并且將前向的輸出 y 作為 user_relu_backward 的輸入 y；將前向的輸出 y 的梯度，作為 user_relu_backward 的輸入 dy；最后輸出 dx。
定完求梯度的方法后，需要調用
ctx->DefineOp(grad_op_name, grad_op_func);
使之生效。
之后的代碼：
const auto& dx_get_func = &ctx, &grad_op_name -> const std::string& {
return ctx->GetOp(grad_op_name).output(“dx”, 0);
};
ctx->FwOp().InputGradBind(user_op::OpArg(“x”, 0), dx_get_func);
是將前向的輸入 x 和剛剛設置的求梯度的方法的輸出(dx) 綁定到一起，這樣，使用 OneFlow 訓練時，就可以自動求導。
其它
? 本文涉及的代碼
? Op 注冊的更多高級設置
? 注冊反向梯度時，也可以使用已有的 Op，而無需專門定制反向 Op。
總結
以上是生活随笔為你收集整理的Python 扩展 Op的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： VS Code 调试 OneFlow
- 下一篇： 自定义 DataLoader