Add support for explicit mode gradients and optimizers

Author: darsnack

Project.toml

@@ -1,8 +1,9 @@
 name = "Flux"
 uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-version = "0.13.7"
+version = "0.13.8-DEV"
 
 [deps]
+AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -26,6 +27,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+AbstractDifferentiation = "0.4.3"
 Adapt = "3.0"
 ArrayInterface = "3.1, 4, 5, 6"
 CUDA = "3"

src/optimise/Optimise.jl

@@ -1,7 +1,15 @@
 module Optimise
 
+using Flux
+using MacroTools: @forward
+import Zygote
+import Zygote: Params, gradient
+using AbstractDifferentiation
+import Optimisers
+import Optimisers: update, update!
 using LinearAlgebra
 import ArrayInterface
+using ProgressLogging: @progress, @withprogress, @logprogress
 
 export train!, update!,
 	Descent, Adam, Momentum, Nesterov, RMSProp,
@@ -10,6 +18,7 @@ export train!, update!,
 	ClipValue, ClipNorm
 
 include("optimisers.jl")
+include("gradients.jl")
 include("train.jl")
 
 end

src/optimise/gradients.jl

@@ -0,0 +1,23 @@
+struct ZygoteImplicitBackend{T} <: AD.AbstractReverseMode
+    core_backend::T
+end
+ZygoteImplicitBackend() = ZygoteImplicitBackend(AD.ZygoteBackend())
+
+AD.@primitive pullback_function(ad::ZygoteImplicitBackend, f, x::Zygote.Params) =
+    AD.pullback_function(ad.core_backend, f, x)
+
+# this is a hack to get around
+# https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/63#issuecomment-1225959150
+AD.gradient(::ZygoteImplicitBackend, f, x::Zygote.Params) = Zygote.gradient(f, x)
+
+struct ZygoteExplicitBackend{T} <: AD.AbstractReverseMode
+    core_backend::T
+end
+ZygoteExplicitBackend() = ZygoteExplicitBackend(AD.ZygoteBackend())
+
+AD.@primitive pullback_function(ad::ZygoteExplicitBackend, f, xs...) =
+    AD.pullback_function(ad.core_backend, f, xs...)
+
+# this is a hack to get around
+# https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/63#issuecomment-1225959150
+AD.gradient(::ZygoteExplicitBackend, f, xs...) = Zygote.gradient(f, xs...)

src/optimise/optimisers.jl

@@ -1,6 +1,3 @@
-using Flux
-using MacroTools: @forward
-
 abstract type AbstractOptimiser end
 
 const EPS = 1e-8

src/optimise/train.jl

@@ -1,7 +1,3 @@
-using ProgressLogging: @progress, @withprogress, @logprogress
-import Zygote: Params, gradient, withgradient
-
-
 """
     update!(opt, p, g)
     update!(opt, ps::Params, gs)
@@ -12,18 +8,23 @@ according to optimizer `opt`  and the gradients `gs` (the gradient `g`).
 As a result, the parameters are mutated and the optimizer's internal state may change.
 The gradient could be mutated as well.
 """
-function update!(opt::AbstractOptimiser, x, x̄)
+function Optimisers.update!(opt::AbstractOptimiser, x, x̄)
   x̄r = ArrayInterface.restructure(x, x̄) # address some cases where Zygote's
                                           # output are not mutable, see #1510 
   x .-= apply!(opt, x, x̄r)
+
+  return opt, x
 end
 
-function update!(opt::AbstractOptimiser, xs::Params, gs)
+function Optimisers.update!(opt::AbstractOptimiser, xs::Params, gs)
   for x in xs
     isnothing(gs[x]) && continue
     update!(opt, x, gs[x])
   end
+
+  return opt, xs
 end
+Optimisers.update(opt::AbstractOptimiser, xs::Params, gs) = update!(opt, xs, gs)
 
 # Callback niceties
 call(f, xs...) = f(xs...)
@@ -82,6 +83,16 @@ end
 batchmemaybe(x) = tuple(x)
 batchmemaybe(x::Tuple) = x
 
+_build_loss(::AD.AbstractBackend, loss, data) = function _loss(m)
+  return loss(m, data...)
+end
+_build_loss(::ZygoteImplicitBackend, loss, data) = function _loss()
+  return loss(data...)
+end
+_gradient_only(x::Zygote.Grads) = x
+_gradient_only(x::NTuple{1}) = x[1]
+_gradient_only(x) = error("Expected gradient w.r.t. single argument (or Zygote.Grads) but got $x")
+
 """
     train!(loss, pars::Params, data, opt::AbstractOptimiser; [cb])
         
@@ -122,19 +133,18 @@ The callback can call [`Flux.stop`](@ref) to interrupt the training loop.
 
 Multiple callbacks can be passed to `cb` as array.
 """
-function train!(loss, ps::Params, data, opt::AbstractOptimiser; cb = () -> ())
+function train!(loss, ad::AD.AbstractBackend, model, data, optstate; cb = () -> ())
   cb = runall(cb)
   itrsz = Base.IteratorSize(typeof(data))
   n = (itrsz == Base.HasLength()) || (itrsz == Base.HasShape{1}()) ? length(data) : 0
   @withprogress for (i, d) in enumerate(data)
     try
-      l, gs = withgradient(ps) do
-        loss(batchmemaybe(d)...)
-      end
+      _loss = _build_loss(ad, loss, batchmemaybe(d))
+      l, gs = AD.valud_and_gradient(ad, _loss, model)
       if !isfinite(l)
         throw(DomainError("Loss is $l on data item $i, stopping training"))
       end
-      update!(opt, ps, gs)
+      optstate, model = update(optstate, model, _gradient_only(gs))
       cb()
     catch ex
       if ex isa StopException
@@ -147,7 +157,11 @@ function train!(loss, ps::Params, data, opt::AbstractOptimiser; cb = () -> ())
     end
     @logprogress iszero(n) ? nothing : i / n
   end
+
+  return optstate, model
 end
+train!(loss, model, data, optstate; kwargs...) =
+  train!(loss, ZygoteImplicitBackend(), model, data, optstate; kwargs...)
 
 """
     @epochs N body