JuliaSmoothOptimizers / JuliaSmoothOptimizers/FluxNLPModels.jl

Speed up the grad

Open
#28 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
7
Forks
4
PR merge metrics
No merged PRs in 30d

Description

I noticed we can improve our grad and objgrad method using the following code
```julia
# Calculate the gradient of the objective
# with respect to the parameters within the model:
grads = Flux.gradient(model) do m
result = m(input)
loss(result, label)
end

```
Currently we have :
```julia
function NLPModels.objgrad!(
nlp::AbstractFluxNLPModel{T, S},
w::AbstractVector{V},
g::AbstractVector{V},
) where {T, S, V}
@lencheck nlp.meta.nvar w g
x, y = nlp.current_training_minibatch

if (eltype(nlp.w) != V) # we check if the types are the same,
update_type!(nlp, w)
g = V.(g)
if eltype(x) != V
x = V.(x)
end
end

increment!(nlp, :neval_obj)
increment!(nlp, :neval_grad)
set_vars!(nlp, w)

f_w = nlp.loss_f(nlp.chain(x), y)
g .= gradient(w_g -> local_loss(nlp, x, y, w_g), w)[1]

return f_w, g
```

We could use

- withgradient(f, args...) for objgrad
- gradient(f, args...) for grad function

This way we won't need local objective call,

reference:
[https://fluxml.ai/Flux.jl/stable/training/zygote/#Zygote.withgradient-Tuple{Any,%20Vararg{Any}}](https://fluxml.ai/Flux.jl/stable/training/zygote/#Zygote.withgradient-Tuple%7BAny,%20Vararg%7BAny%7D%7D)

Note: I will be doing this as a pr soon but thought it would be nice to have a guide here just in case

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.