写算子单元测试Writing Unit Tests!

2023-12-25 15:21:46

一些单元测试示例，可在tests/python/relay/test_op_level3.py中找到，用于累积总和与乘积算子。

梯度算子

梯度算子对于编写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和要累加的梯度。返回的是一个列表，第i个索引处的元素是算子相对于算子第i个输入的导数。通常，梯度将返回一个列表，包含的元素数量与基本算子的输入数量相同。

进一步分析这个定义前面，首先回顾一下sigmoid函数的导数：

上面的定义类似于数学定义，但有一个重要的补充，将在下面描述。

术语orig*（类似于（orig）-orig）直接匹配导数，这里的orig是sigmoid函数，但不只是对如何计算这个函数的梯度感兴趣。将这个梯度与其它梯度组合起来，这样就可在整个程序中累积梯度。

这就是梯度术语的意义所在。在表达式grad*orig*（one_like（orig）-orig）中，乘以grad，表示如何使用到目前为止的梯度合成导数。

现在，考虑乘法，一个稍微有趣的示例：

@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是一个二进制算子。回想一下，如果

，偏导数是

有一个乘法所需的步骤，对于sigmoid不是必需的，乘法具有广播语义。梯度的shape可能与输入的shape不匹配，使用collapse_sum_like来获取梯度grad * <var>项的内容，使shape与要区分的输入的shape匹配。

在C++中添加一个梯度，类似于在Python中添加，但是，用于注册的接口略有不同。

首先，确保包含src/relay/transforms/pattern_utils.h。提供了用于在RelayAST中创建节点的 helper函数。然后，类似于Python示例的方式，定义梯度：

tvm::Array<Expr> MultiplyGrad(const Expr& orig_call，const Expr& output_grad) {

const Call& call=orig_call.Downcast<Call>();

return { CollapseSumLike(Multiply(output_grad，call.args[1])，call.args[0])，

CollapseSumLike(Multiply(output_grad，call.args[0])，call.args[1]) };

}

在C++中，不能使用Python中的算子重载，需要进行downcast，实现更加冗长。即使如此，可容易地验证这个定义，是否反映了Python中的早期示例。

现在，不需要使用Python装饰器，需要在基础算子的注册末尾，添加一个对“FPrimalGradient”的set_attr调用，注册梯度。

RELAY_REGISTER_OP("multiply")

    // ...

    // Set other attributes

    // ...

    .set_attr<FPrimalGradient>("FPrimalGradient"，MultiplyGrad);

参考链接：

https://tvm.apache.org/docs/dev/relay_add_op.html

TVM源码研习 — TVM中的IR设计与技术实现

码农公寓