自动微分

神经网络的训练使用反向传播算法求梯度。仓颉 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 的输入是 xy,且两个输入都可微:

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 同类型,对应了 xy 的梯度。
  • 函数 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)
    }
}