利用合子计算Julia中包含NN w.r.t参数的损失函数的Hessian

如何计算由神经网络w.r.t.神经网络参数组成的损失函数的hessian值

例如，考虑

以下的损失函数

using Flux: Chain, Dense, σ, crossentropy, params
using Zygote
model = Chain(
    x -> reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -> σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss = model -> crossentropy(model(input), target)

我可以用两种方法得到梯度w.r.t参数

Zygote.gradient(model -> loss(model), model)

或

然而，我找不到一种方法来获取hessian w.r.t及其参数。（我想做一些类似于

Zygote.hessian（model->loss（model），model）

，但是

Zygote.hessian

不接受

：：Params

作为输入）

最近，一个

jacobian

函数被发送到主分支（发行号#910），该分支

我一直在尝试结合

梯度

和

雅可比矩阵

来得到一个海森矩阵（因为海森矩阵是函数梯度的雅可比矩阵），但是没有用。 我认为问题在于

模型

是一个

链

对象，它包括

重塑

和

σ等通用函数。

缺少参数，但我无法克服这一点

grad = model -> Zygote.gradient(model -> loss(model), model)
jacob = model -> Zygote.jacobian(grad, model)
jacob(model) ## does not work

编辑：作为参考，我已经

grad = model -> Zygote.gradient(model -> loss(model), model)
jacob = model -> Zygote.jacobian(grad, model)
jacob(model) ## does not work