Julia 自定义通量梯度，而不是使用受精卵A.D

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

W = rand(Nh, N)
U = rand(N, Nh)
b = rand(N)
c = rand(Nh)

θ = Flux.Params([b, c, U, W])

opt = ADAM(0.01)

然后我有一个函数来计算模型参数的解析梯度，

θ

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

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

grads = gradients(x)
update!(opt, θ, grads)

---

我的问题是：我的

gradient（x）

函数需要返回什么形式/类型才能执行

更新！（opt，θ，grads）

，我该怎么做？

如果不使用

Params

，那么

grads

只需要是梯度。唯一的要求是

θ

和

梯度

大小相同

例如，

map（（x，g）->update！（opt，x，g），θ，grads）

其中

θ==[b，c，U，W]

和

grads=[gradients（b），gradients（c），gradients（U），gradients（W）]

（不确定

gradients

需要什么样的输入）

更新：但要回答您的原始问题，

gradients

需要返回一个

Grads

对象，该对象位于以下位置：

大概是

# within gradient function body assuming gb is the gradient w.r.t b
g = Zygote.Grads(IdDict())
g.grads[θ[1]] = gb # assuming θ[1] == b

但是不使用

Params

可能更容易调试。唯一的问题是没有

更新将处理一系列参数，但您可以轻松定义自己的参数：

function Flux.Optimise.update!(opt, xs::Tuple, gs)
    for (x, g) in zip(xs, gs)
        update!(opt, x, g)
    end
end

# use it like this
W = rand(Nh, N)
U = rand(N, Nh)
b = rand(N)
c = rand(Nh)

θ = (b, c, U, W)

opt = ADAM(0.01)
x = # generate input to gradients
grads = gradients(x) # return tuple (gb, gc, gU, gW)
update!(opt, θ, grads)

更新2：

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

W = rand(Nh, N)
U = rand(N, Nh)
b = rand(N)
c = rand(Nh)

θ = Flux.Params([b, c, U, W])

f(x) = # define your model
gradients(x) = # define your analytical gradient

# set up the custom adjoint
Zygote.@adjoint f(x) = f(x), Δ -> (gradients(x),)

opt = ADAM(0.01)
x = # generate input to model
y = # output of model
grads = Zygote.gradient(() -> Flux.mse(f(x), y), θ)
update!(opt, θ, grads)

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

如果您正在做更复杂的事情，并且不能使用
Zygote.gradient
，那么第一种方法（不使用
Params
）最合适
Params实际上只存在于与Flux的旧广告向后兼容的情况下，因此如果可能的话，最好避免使用它。
我使用了你的第一个建议，效果很好。非常感谢！