写算子单元测试Writing Unit Tests! 一些单元测试示例,可在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函数,但不只是对如何计算这个函数的梯度感兴趣。将这个梯度与其它梯度组合起来,这样就可在整个程序中累积梯度。
这就是梯度术语的意义所在。在表达式gradorig(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 * 项的内容,使shape与要区分的输入的shape匹配。 Adding a Gradient in C++ 在C++中添加一个梯度,类似于在Python中添加,但是,用于注册的接口略有不同。 首先,确保包含src/relay/transforms/pattern_utils.h。提供了用于在RelayAST中创建节点的 helper函数。然后,类似于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中的算子重载,需要进行downcast,实现更加冗长。即使如此,可容易地验证这个定义,是否反映了Python中的早期示例。 现在,不需要使用Python装饰器,需要在基础算子的注册末尾,添加一个对“FPrimalGradient”的set_attr调用,注册梯度。 RELAY_REGISTER_OP(“multiply”) // … // Set other attributes // … .set_attr(“FPrimalGradient”,MultiplyGrad);
参考链接: https://tvm.apache.org/docs/dev/relay_add_op.html TVM源码研习 — TVM中的IR设计与技术实现
|