向Relay添加算子
為了在Relay IR中使用TVM算子，需要在Relay中注冊算子，以確保將其集成到Relay的類型系統中。
注冊算子需要三個步驟：
? 使用RELAY_REGISTER_OPC ++中的宏注冊算子的Arity和類型信息
? 定義一個C ++函數為算子生成一個調用節點，并為該函數注冊一個Python API掛鉤
? 將上述Python API掛鉤包裝在更整潔的界面中
該文件src/relay/op/tensor/binary.cc提供了前兩個步驟的python/tvm/relay/op/tensor.py示例，同時提供了后兩個步驟的示例。
注冊算子
TVM已經具有算子注冊表，如果沒有其它類型信息，Relay無法正確合并TVM算子。
為了在注冊算子時具有更大的靈活性，并在Relay中表達類型時，提高了表達性和粒度，使用輸入和輸出類型之間的關系來對算子進行類型化。這些關系表示為接受輸入類型和輸出類型的列表（這些類型中的任何一個都不完整），返回滿足該關系的輸入和輸出類型的列表的函數。本質上，算子的關系除了計算輸出類型外，還可以強制執行所有必要的鍵入規則（即，通過檢查輸入類型）。
例如，參閱src/relay/op/type_relations.h及其實現。例如，BroadcastRel接受兩個輸入類型和一個輸出類型，都是具有相同基礎數據類型的張量類型，最后確保輸出類型的形狀是輸入類型的形狀的廣播。
type_relations.h 如果現有類型未捕獲所需算子的行為，則可能有必要添加另一種類型關系。
使用RELAY_REGISTER_OPC ++中的宏，開發人員可以在Relay中指定有關算子的以下信息：
? Arity（參數個數）
? 位置參數的名稱和說明
? 支持級別（1表示內部固有；較高的數字表示積分較少或不受外部支持的算子）
? 算子的類型關系
下面的示例來自binary.cc張量，并將其用于張量廣播：
RELAY_REGISTER_OP(“add”)
.set_num_inputs(2)
.add_argument(“lhs”, “Tensor”, “The left hand side tensor.”)
.add_argument(“rhs”, “Tensor”, “The right hand side tensor.”)
.set_support_level(1)
.add_type_rel(“Broadcast”, BroadcastRel);
創建調用節點
此步驟僅需要簡單地編寫一個將參數帶給算子的函數（如Relay表達式），然后將調用節點返回給算子（即，應將其放置在Relay AST中的節點，該AST是要向算子進行調用的位置）。
目前不支持調用屬性和類型參數（最后兩個字段），足以用于Op::Get從算子注冊表中獲取算子信息，并將參數傳遞給調用節點，如下所示。
TVM_REGISTER_GLOBAL(“relay.op._make.add”)
.set_body_typed<Expr(Expr, Expr)>([](Expr lhs, Expr rhs) {
static const Op& op = Op::Get(“add”);
return Call(op, {lhs, rhs}, Attrs(), {});
});
包括Python API掛鉤
在Relay中通常是約定，通過導出的函數TVM_REGISTER_GLOBAL應該包裝在單獨的Python函數中，而不是在Python中直接調用。對于產生對算子的調用的函數，捆綁起來可能很方便，其中python/tvm/relay/op/tensor.py，都提供了張量上的元素算子。例如，以下是上一節中的add函數在Python中的顯示方式：
def add(lhs, rhs):
“”"Elementwise addition.

Parameters
----------
lhs : relay.ExprThe left hand side input data
rhs : relay.ExprThe right hand side input dataReturns
-------
result : relay.ExprThe computed result.
"""
return _make.add(lhs, rhs)

這些Python庫也可能是向算子提供更簡單接口的好機會。例如，該 concat算子被注冊為僅接受一個算子，即一個具有要連接的張量的元組，Python庫將這些張量作為參數，組合成一個元組，然后生成調用節點：
def concat(*args):
“”"Concatenate the input tensors along the zero axis.

Parameters
----------
args: list of TensorReturns
-------
tensor: The concatenated tensor.
"""
tup = Tuple(list(args))
return _make.concat(tup)

梯度算子
梯度算子對于在Relay中編寫可區分的程序很重要。盡管Relay的autodiff算法可以區分一流的語言結構，算子是不透明的。由于Relay無法調查實現，必須提供明確的區分規則。
Python和C ++均可用于編寫梯度算子，將示例重點放在Python上，因為更常用。
在Python中添加漸變
可以找到Python梯度算子的集合 python/tvm/relay/op/_tensor_grad.py。將通過兩個有代表性的示例：sigmoid和multiply。
@register_gradient(“sigmoid”)
def sigmoid_grad(orig, grad):
“”“Returns [grad * sigmoid(x) * (1 - sigmoid(x))].”""
return [grad * orig * (ones_like(orig) - orig)]
這里的輸入是原始算子orig和grad要累加到的漸變。返回的是一個列表，其中第i個索引處的元素是算子相對于算子第i個輸入的派生。通常，漸變將返回一個列表，其中包含與基本算子輸入相同數量的元素。
在進一步分析這個定義之前，應該回想一下S型函數的導數： ?σ?x=σ(x)(1?σ(x))?σ?x=σ(x)(1?σ(x))。上面的定義看起來與數學定義相似，有一個重要的補充，將在下面進行描述。
該術語直接與導數匹配，因為這是S型函數，不僅對如何計算此函數的梯度感興趣。有興趣將此梯度與其它梯度組成，可以在整個程序中累積該梯度。這是該術語出現的地方。指定到目前為止如何用梯度組成導數。orig * (ones_like(orig) - orig)origgradgrad * orig * (ones_like(orig) - orig)grad
來看multiply一個更有趣的示例：
@register_gradient(“multiply”)
def multiply_grad(orig, grad):
“”“Returns [grad * y, grad * x]”""
x, y = orig.args
return [collapse_sum_like(grad * y, x),
collapse_sum_like(grad * x, y)]
在此示例中，返回列表中有兩個元素， multiply是一個二進制算子。如果f(x,y)=xyf(x,y)=xy，偏導數是 ?f?x=y?f?x=y 和 ?f?y=x?f?y=x。
有一個必需的步驟，因為廣播具有語義，multiply不需要。可能與輸入的形狀不匹配，習慣于獲取術語的內容，并使形狀與要區分的輸入的形狀相匹配。sigmoidmultiplygradcollapse_sum_likegrad *
在C ++中添加漸變
在C ++中添加漸變類似于在Python中添加漸變，注冊的界面略有不同。
確保src/relay/pass/pattern_utils.h包含其中。提供了用于在Relay AST中創建節點的輔助功能。然后，以類似于Python示例的方式定義漸變：
tvm::Array MultiplyGrad(const Expr& orig_call, const Expr& output_grad) {
const Call& call = orig_call.Downcast();
return { CollapseSumLike(Multiply(output_grad, call.args[1]), call.args[0]),
CollapseSumLike(Multiply(output_grad, call.args[0]), call.args[1]) };
}
在C ++中，不能使用與Python中相同的算子重載，并且需要向下轉換，實現更為冗長。即使這樣，仍可以輕松地驗證此定義，是否與Python中的先前示例相同。
代替使用Python渲染，需要set_attr在基本算子注冊的末尾添加對“ FPrimalGradient”的調用，以注冊漸變。
RELAY_REGISTER_OP(“multiply”)
// …
// Set other attributes
// …
.set_attr(“FPrimalGradient”, MultiplyGrad);
總結
? TVM算子可以使用表示適當類型信息的關系在中繼中注冊。
? 在中繼中使用算子需要一個函數來為算子生成調用節點。
? 最好有一個簡單的Python庫來生成調用節點。

總結

以上是生活随笔為你收集整理的向Relay添加算子的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。