# How to get the results and gradients when using ForwardDiff.jl

**URL:** <https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525>\
**Category:** Machine Learning\
**Created:** [November 20, 2022, 6:23am UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525 "2022-11-20T06:23:15Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![Frankiewaang](https://avatars.discourse-cdn.com/v4/letter/f/f9ae1b/32.png) [@Frankiewaang](https://discourse.julialang.org/u/Frankiewaang)\
**Post date:** [November 20, 2022, 6:23am UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/1 "2022-11-20T06:23:15Z")

</div>

I did some research and found out that if I want to get both the value of a function and its gradients I could use the API from `DiffResults`. But if the result from a function returns a tuple(instead of a scalar) where the first one is the loss but I also want to get the second and third results how should I do it?

In `Zygote.jl`, I could do something like:

```julia
#loss function for Lux model
function core(S, dcf, ps, model, st)
    value_if_wait, st = model(S, ps, st)
    mse = mean(abs2, value_if_wait .- dcf)
    return mse, value_if_wait, st
end
(loss, value_if_wait, st), back = pullback(p -> core(S_i, dcf, p, model, st), ps)
gs = back((one(loss), nothing, nothing))[1]
#and then do the update 

```

---

<div class="post-metadata">

**Author:** ![colleybrb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/colleybrb/32/44529_2.png) [@colleybrb](https://discourse.julialang.org/u/colleybrb)\
**Post date:** [November 21, 2022, 2:42pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/2 "2022-11-21T14:42:45Z")

</div>

Not being familair with pkg… I’d see what it’s currently returning by looking in the source code, or return under on variable and print. If it is returning you should be able to match the return to the variable assignment. If not you will need to add additional returns.

Hope that helps.

---

<div class="post-metadata">

**Author:** ![Frankiewaang](https://avatars.discourse-cdn.com/v4/letter/f/f9ae1b/32.png) [@Frankiewaang](https://discourse.julialang.org/u/Frankiewaang)\
**Post date:** [November 21, 2022, 4:19pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/3 "2022-11-21T16:19:09Z")

</div>

Not worries. I just saw the slack discussion [here](https://julialang.slack.com/archives/C690QRAA3/p1668765629120249) and wanna give it a try, my thinking is it seems to be easier for `Lux` to work with other AD libraries and I only need to flatten it to call the `gradient` function and reconstruct it when I wanna call the model. For Flux, the params are stored in the struct, therefore a customized walk is needed.

---

<div class="post-metadata">

**Author:** ![j-fu](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/j-fu/32/11373_2.png) [@j-fu](https://discourse.julialang.org/u/j-fu)\
**Post date:** [November 21, 2022, 4:23pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/4 "2022-11-21T16:23:23Z")

</div>

Assume you have mutating function `func!(y,u)` which writes the result into `y`.  
Prepare a result buffer (once):

```julia
diffresult=DiffResults.JacobianResult(u0)
y=zero(u0)
cfg = ForwardDiff.JacobianConfig(func!,y, u0) 

```

Then you should be able to call (many times)

```julia
ForwardDiff.jacobian!(diffresult, func!,y,u,cfg)

```

and access `DiffResults.value(diffresult)` and `DiffResults.jacobian(diffresult)` without allocations.

Not sure if there is another way for your case, but I think you can wrap this around your `core())` function.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 21, 2022, 4:27pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/5 "2022-11-21T16:27:28Z")

</div>

It is easy to use a Flux model with a flat vector of parameters, the [Optimisers docs](https://fluxml.ai/Optimisers.jl/dev/#Obtaining-a-flat-parameter-vector) have an example using ForwardDiff, and the [Flux docs](https://fluxml.ai/Flux.jl/stable/destructure/) have one with a Hessian.

---

<div class="post-metadata">

**Author:** ![lazarusA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lazarusa/32/6571_2.png) [@lazarusA](https://discourse.julialang.org/u/lazarusA)\
**Post date:** [November 21, 2022, 9:50pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/6 "2022-11-21T21:50:22Z")

</div>

@mcabbott Because I was looking for something with ForwardDiff I tested the example from Optimisers and compared the output with Zygote, unfortunately they are not the same 😢 , take a look at the loss output. Am I doing something wrong?

Here both versions.

- With ForwardDiff and Optimisers

```nohighlight
using ForwardDiff # an example of a package which only likes one array
using Flux
using Random
using Optimisers
Random.seed!(123)

model = Chain( # much smaller model example, as ForwardDiff is a slow algorithm here
          Conv((3, 3), 3 => 5, pad=1, bias=false), 
          BatchNorm(5, relu), 
          Conv((3, 3), 5 => 3, stride=16),
        )
image = rand(Float32, 224, 224, 3, 1);
@show sum(model(image));

loss(m, x) = sum(m(x))

rule = Optimisers.Adam(0.001f0, (0.9f0, 0.999f0), 1.1920929f-7)

flat, re = Flux.destructure(model)
st = Optimisers.setup(rule, flat) # state is just one Leaf now

∇flat = ForwardDiff.gradient(flat) do v
    loss(re(v), image) # re(v), rebuild a new object like model
end

st, flat = Optimisers.update(st, flat, ∇flat)
@show loss(re(flat),image);

```

```julia
sum(model(image)) = -0.33076355f0
loss(re(flat), image) = -7.7023053f0

```

- And here the one with Zygote.

```nohighlight
using Flux
using Random
Random.seed!(123)

model = Chain( # much smaller model example, as ForwardDiff is a slow algorithm here
          Conv((3, 3), 3 => 5, pad=1, bias=false), 
          BatchNorm(5, relu), 
          Conv((3, 3), 5 => 3, stride=16),
        )
image = rand(Float32, 224, 224, 3, 1);
@show sum(model(image));

loss(m, x) = sum(m(x))

opt = Flux.Adam(0.001f0, (0.9f0, 0.999f0), 1.1920929f-7)
θ = Flux.params(model)
grads = Flux.gradient(θ) do 
    loss(model, image)
end

Flux.update!(opt, θ, grads)
@show loss(model, image);

```

with this

```julia
sum(model(image)) = -0.33076355f0
loss(model, image) = -5.064876f0

```

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 21, 2022, 9:57pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/7 "2022-11-21T21:57:20Z")

</div>

That’s no good, can you make an issue?

I think the core is that `BatchNorm` has a test/train-mode change, which doesn’t happen with ForwardDiff. Commenting out that layer leads to identical results.

---

<div class="post-metadata">

**Author:** ![lazarusA](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lazarusa/32/6571_2.png) [@lazarusA](https://discourse.julialang.org/u/lazarusA)\
**Post date:** [November 21, 2022, 10:16pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/8 "2022-11-21T22:16:00Z")

</div>

> [@lazarusA](#):
>
> ested the example from Optimisers and compared the output with Zygote, unfortunately they are not the same

I see. Ok, here the report: [ForwardDiff + destructure is different from Zygote, on a model with BatchNorm · Issue #2122 · FluxML/Flux.jl · GitHub](https://github.com/FluxML/Optimisers.jl/issues/117)

---

<div class="post-metadata">

**Author:** ![tchebycheff](https://avatars.discourse-cdn.com/v4/letter/t/779978/32.png) [@tchebycheff](https://discourse.julialang.org/u/tchebycheff)\
**Post date:** [November 23, 2022, 1:49am UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/9 "2022-11-23T01:49:50Z")

</div>

This seems to be specific to this particular loss function and/or model.

Consider this example otoh:

```julia
using Flux, ForwardDiff, Random
Random.seed!(123)
mlp = Chain(Dense(20, 16, relu), Dense(16,8,relu), Dense(8,1,σ))
ps, re = Flux.destructure(mlp) 
xs = randn(20,50)
ys = mapslices(x->exp.(sin.(sum(x))), xs, dims=1)

bar(p) = Flux.mse(re(p)(xs), ys)

d1 = ForwardDiff.gradient(bar, ps)
d2 = Flux.gradient(bar, ps)[1]

d1 ≈ d2 # true

```

---

<div class="post-metadata">

**Author:** ![colleybrb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/colleybrb/32/44529_2.png) [@colleybrb](https://discourse.julialang.org/u/colleybrb)\
**Post date:** [November 25, 2022, 4:21pm UTC](https://discourse.julialang.org/t/how-to-get-the-results-and-gradients-when-using-forwarddiff-jl/90525/10 "2022-11-25T16:21:24Z")

</div>

@edit macro to get to the source code next time and if you are walking a struct I’ve found ComponentArrays as a good alternative to structs. Hope that helps for future.