Flux 的自定义梯度而不是使用 Zygote AD

Question

我有一个机器学习模型，其中模型参数的梯度是解析的，不需要自动微分。然而，我仍然希望能够利用 Flux 中的不同优化器，而不必依赖 Zygote 进行区分。这是我的代码的一些片段。

\n\n

W = rand(Nh, N)\nU = rand(N, Nh)\nb = rand(N)\nc = rand(Nh)\n\n\xce\xb8 = Flux.Params([b, c, U, W])\n\nopt = ADAM(0.01)\n

\n\n

然后我有一个函数可以计算模型参数的解析梯度，\xce\xb8。

\n\n

function gradients(x) # x = one input data point or a batch of input data points\n    # stuff to calculate gradients of each parameter\n    # returns gradients of each parameter\n

\n\n

然后我希望能够做如下的事情。

\n\n

grads = gradients(x)\nupdate!(opt, \xce\xb8, grads)\n

\n\n

我的问题是：我的函数需要返回什么形式/类型gradient(x)才能执行此操作update!(opt, \xce\xb8, grads)，以及如何执行此操作？

\n

Answer 1

如果你不使用Params那么grads只需要渐变。唯一的要求是\xce\xb8和grads大小相同。

\n\n

例如，map((x, g) -> update!(opt, x, g), \xce\xb8, grads)哪里\xce\xb8 == [b, c, U, W]和grads = [gradients(b), gradients(c), gradients(U), gradients(W)]（不太确定gradients您期望输入什么）。

\n\n

更新：但要回答您原来的问题，gradients需要返回Grads在这里找到的对象： https: //github.com/FluxML/Zygote.jl/blob/359e586766129878ca0e56121037ed80afda6289/src/compiler/interface.jl#L88

\n\n

所以像

\n\n

# within gradient function body assuming gb is the gradient w.r.t b\ng = Zygote.Grads(IdDict())\ng.grads[\xce\xb8[1]] = gb # assuming \xce\xb8[1] == b\n

\n\n

但不使用Params可能更容易调试。唯一的问题是没有update!可以处理参数数组的方法，但您可以轻松定义自己的参数：

\n\n

function Flux.Optimise.update!(opt, xs::Tuple, gs)\n    for (x, g) in zip(xs, gs)\n        update!(opt, x, g)\n    end\nend\n\n# use it like this\nW = rand(Nh, N)\nU = rand(N, Nh)\nb = rand(N)\nc = rand(Nh)\n\n\xce\xb8 = (b, c, U, W)\n\nopt = ADAM(0.01)\nx = # generate input to gradients\ngrads = gradients(x) # return tuple (gb, gc, gU, gW)\nupdate!(opt, \xce\xb8, grads)\n

\n\n

更新2：

\n\n

另一种选择是仍然使用 Zygote 获取梯度，以便它自动Grads为您设置对象，但使用自定义伴随，以便它使用您的分析函数来计算伴随。假设您的 ML 模型定义为名为 的函数f，以便f(x)返回输入的模型输出x。我们还假设gradients(x)返回分析梯度，x就像您在问题中提到的那样。然后以下代码仍将使用 Zygote 的 AD 来Grads正确填充对象，但它将使用您为函数计算梯度的定义f：

\n\n

W = rand(Nh, N)\nU = rand(N, Nh)\nb = rand(N)\nc = rand(Nh)\n\n\xce\xb8 = Flux.Params([b, c, U, W])\n\nf(x) = # define your model\ngradients(x) = # define your analytical gradient\n\n# set up the custom adjoint\nZygote.@adjoint f(x) = f(x), \xce\x94 -> (gradients(x),)\n\nopt = ADAM(0.01)\nx = # generate input to model\ny = # output of model\ngrads = Zygote.gradient(() -> Flux.mse(f(x), y), \xce\xb8)\nupdate!(opt, \xce\xb8, grads)\n

\n\n

请注意，我Flux.mse在上面使用了损失作为示例。这种方法的一个缺点是 Zygote 的gradient函数需要标量输出。如果您的模型被传递到一些将输出标量误差值的损失，那么这@adjoint是最好的方法。这适用于您正在进行标准 ML 的情况，唯一的变化是您希望 Zygote 计算f使用您的函数来计算分析梯度。

\n\n

如果您正在做更复杂的事情并且无法使用Zygote.gradient，那么第一种方法（不使用Params）是最合适的。Params实际上只是为了向后兼容 Flux 的旧 AD，所以如果可能的话最好避免它。

\n