自动微分
神经网络的训练使用反向传播算法求梯度。仓颉 TensorBoost 提供函数式自动微分的能力,应用链式求导法则,自动微分对用户屏蔽了大量的求导细节和过程。一般而言,用户可对被 @Differentiable 修饰的任一函数进行自动求导,通过微分表达式计算函数入参对应的梯度。用户还可通过微分表达式组合进行高阶微分,完成更强大的求导功能。
可微函数和微分表达式
仓颉 TensorBoost 微分表达式作用于函数,对函数入参求取梯度。以网络模型训练过程为例:
@Differentiable
首先定义训练函数 train
,使用 @Differentiable 修饰函数,include 列表包含 train 函数第一个入参(定制网络结构 MyLayer)。
@Differentiable[include: [net]]
func train(net: MyLayer, input: Tensor, label: Tensor): Tensor {
let out = net(input) // 执行 MyLayer operator() 方法
(out - label) ** 2.0
}
@AdjointOf
要对训练函数进行求导,则需要用到 @AdjointOf
表达式:
let net = MyLayer(16, 16, true)
let input = onesTensor(Array<Int64>([16, 16]))
let label = onesTensor(Array<Int64>([1]))
let adj = @AdjointOf(train)
let (out, bpFunc) = adj(net, input, label)
let netGrad: MyLayer = bpFunc(onesLike(out))
@AdjointOf
为仓颉 TensorBoost 提供的微分表达式。下面的代码是 @AdjointOf
的使用样例,假设函数 f
的输入是 x
和 y
,且两个输入都可微:
let adj = @AdjointOf(f)
let (out, bpFunc) = adj(x, y)
let sensitivity = out
let (dx, dy) = bpFunc(sensitivity)
【说明】:
@AdjointOf
的返回值为原函数f
的伴随函数adj
。- 伴随函数
adj
的返回值为原函数f
的正向计算结果out
和微分计算函数bpFunc
。 bpFunc
接受和out
同类型的sensitivity
值,输出(dx, dy)
,dx/dy
分别和x/y
同类型,对应了x
和y
的梯度。- 函数
f
必须为被@Differentiable
修饰的全局的微分函数。
stopGradient
用户在编写可微函数时可能会涉及到一些运算操作不参与或不支持微分求导。如果这些运算操作访问到参与微分的数据,意味着链式求导过程出现中断。此时需要用户使用仓颉 TensorBoost 提供的 stopGradient
函数接口来显示标注需要停止求导的数据。
比如下面代码中,可微方法 SmartMatMul
根据输入 Tensor 维度选择合适的运算方式。获取 Tensor 维度通过 getShapeDims()
。而这个方法不参与微分求导,所以需要对参与微分的 input
先进行 stopGradient()
再调用 getShapeDims()
。
@Differentiable
func SmartMatMul(input: Tensor, weight: Tensor): Tensor {
return if (stopGradient(input).getShapeDims() == 2) {
// 2D Tensor:调用 matmul
matmul(input, weight)
} else {
// 超过 2D Tensor:调用 batchMatMul
batchMatMul(input, weight)
}
}