JuliaSmoothOptimizers / JuliaSmoothOptimizers/FluxNLPModels.jl
Speed up the grad
- 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
Assessment
This issue has not been assessed yet.