search

JuliaNLSolvers / Optim.jl

Optimization functions for Julia
Other
1.11k stars 213 forks source link

Autodiff using Enzyme #1001

Closed renatobellotti closed 1 year ago

renatobellotti commented 1 year ago

Hi,

I'm looking for an optimisation library in Julia that allows to use auto-differentiation using Enzyme. I have seen that there is an option to use autodiff, but need fine-grained control. Namely, I need to propagate the autodiff through custom CPU/GPU kernels because the problem I'm solving is quite computationally expensive. I know that Enzyme can do such a thing, but is it possible to use Enzyme.jl with Optim.jl?

dehann commented 1 year ago

hi, any reason for Enzyme vs Zygote?

antoine-levitt commented 1 year ago

The questions of optimizers and AD are mostly orthogonal. Figure out how to compute gradients using AD, then pass to optim that gradient.

renatobellotti commented 1 year ago

@dehann Yes, as far as I know only Enzyme can calculate gradients for custom kernels.

renatobellotti commented 1 year ago

@antoine-levitt Nice, this sounds very powerful. Can you please elaborate on how to do this? I interpret the documentation [1] in the way that we either have to use finite differences or a predefined AD library. Am I mistaken?

[1] https://julianlsolvers.github.io/Optim.jl/v0.9.3/algo/autodiff/

antoine-levitt commented 1 year ago

Simply write a g(x) function that does whatever_ad(f, x) and pass it to optim as en explicit gradient

renatobellotti commented 1 year ago

Nice, thanks! :)