Add gradient clipping (#27)

* add gradient clipping

* include in OptimiserChain tests

* from review

* docs

* rebase fixup

* wording

Author: mcabbott

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

docs/src/api.md

@@ -13,6 +13,11 @@ Optimisers.AMSGrad
 Optimisers.NADAM
 Optimisers.ADAMW
 Optimisers.AdaBelief
+```
+
+```
+Optimisers.ClipGrad
+Optimisers.ClipNorm
 Optimisers.WeightDecay
 Optimisers.OptimiserChain
 ```

src/Optimisers.jl

@@ -1,12 +1,13 @@
 module Optimisers
 
 using Functors: functor, fmap, isleaf
+using LinearAlgebra
 
 include("interface.jl")
 include("rules.jl")
 
 export Descent, ADAM, Momentum, Nesterov, RMSProp,
        ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
-       WeightDecay, OptimiserChain
+       WeightDecay, ClipGrad, ClipNorm, OptimiserChain
 
 end # module

src/rules.jl

@@ -490,6 +490,61 @@ function apply(o::WeightDecay, state, x, dx)
   return state, dx′
 end
 
+"""
+    ClipGrad(δ = 10f0)
+
+Restricts every gradient component to obey `-δ ≤ dx[i] ≤ δ`.
+
+See also [`ClipNorm`](@ref).
+"""
+struct ClipGrad{T<:Real}
+  delta::T
+end
+ClipGrad() = ClipGrad(10f0)
+
+init(o::ClipGrad, x::AbstractArray) = nothing
+
+(o::ClipGrad)(state::Nothing, m, dm) = update(o, state, m, dm)
+
+function apply(o::ClipGrad, state, x, dx)
+  δ = convert(eltype(dx), o.delta)
+  dx′ = @. clamp(dx, -δ, δ)
+
+  return state, dx′
+end
+
+"""
+    ClipNorm(ω = 10f0, p = 2; throw = true)
+
+Scales any gradient array for which `norm(dx, p) > ω`
+to stay at this threshold (unless `p==0`).
+
+Throws an error if the norm is infinite or `NaN`,
+which you can turn off with `throw = false`.
+
+See also [`ClipGrad`](@ref).
+"""
+struct ClipNorm{T<:Real}
+  omega::T
+  p::T
+  throw::Bool
+end
+ClipNorm(ω = 10f0, p = 2; throw::Bool = true) = ClipNorm{typeof(ω)}(ω, p, throw)
+
+init(o::ClipNorm, x::AbstractArray) = nothing
+
+(o::ClipNorm)(state::Nothing, m, dm) = update(o, state, m, dm)
+
+function apply(o::ClipNorm, state, x, dx)
+  nrm = norm(dx, o.p)
+  if o.throw && !isfinite(nrm)
+    throw(DomainError("gradient has $(o.p)-norm $nrm, for array $(summary(x))"))
+  end
+  λ = min(o.omega / nrm, 1)
+
+  return state, @. dx * λ
+end
+
 """
     OptimiserChain(opts...)
 

test/runtests.jl

@@ -1,13 +1,14 @@
 using Optimisers, Test
-using Zygote, Random
-using Statistics
+using Zygote
+using Statistics, Random, LinearAlgebra
+Random.seed!(84)
+
+@testset verbose=true "Optimisers.jl" begin
 
-@testset "Optimisers" begin
-  Random.seed!(84)
-  w′ = (α = rand(3, 3), β = rand(3, 3))
   @testset for o in (Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
                      ADAGrad(), AdaMax(), ADADelta(), AMSGrad(), NADAM(),
                      ADAMW(), RADAM(), OADAM(), AdaBelief())
+    w′ = (α = rand(3, 3), β = rand(3, 3))
 
     # Original example
     w = (α = 5rand(3, 3), β = rand(3, 3))
@@ -37,25 +38,50 @@ using Statistics
     end
 
   end
-end
 
-@testset "OptimiserChain" begin
-  Random.seed!(84)
-  w = randn(10, 10)
-  w′ = randn(10, 10)
-  loss(x, w, w′) = mean((w*x .- w′*x) .^ 2)
-  opt = OptimiserChain(WeightDecay(), ADAM(0.001))
-  st = Optimisers.state(opt, w)
-  for t = 1:10^5
-    x = rand(10)
-    gs = gradient(w -> loss(x, w, w′), w)
-    st, w = Optimisers.update(opt, st, w, gs...)
+  @testset "OptimiserChain with $pre" for pre in (WeightDecay(), ClipGrad(), ClipNorm())
+    Random.seed!(84)
+    w = randn(10, 10)
+    w′ = randn(10, 10)
+    loss(x, w, w′) = mean((w*x .- w′*x) .^ 2)
+    @test loss(rand(10, 10), w, w′) > 1
+    opt = OptimiserChain(pre, ADAM(0.001))
+    st = Optimisers.init(opt, w)
+    for t = 1:10^5
+      x = rand(10)
+      gs = gradient(w -> loss(x, w, w′), w)
+      st, w = Optimisers.update(opt, st, w, gs...)
+    end
+    @test loss(rand(10, 10), w, w′) < 0.01
+  end
+
+  @testset "gradient clipping" begin
+    @test_skip m = (α = ([0], sin), γ = rand(3))  # https://github.com/FluxML/Optimisers.jl/issues/28
+    m = (α = ([0], [0]), γ = rand(3))
+    c1 = ClipGrad(13)
+    s1 = Optimisers.state(c1, m)
+    _, g1 = Optimisers.update(c1, s1, m, (α = nothing, γ = [1,10,100],))
+    @test m.γ .- g1.γ ≈ [1, 10, 13]
+
+    c2 = ClipNorm(10)
+    s2 = Optimisers.state(c2, m)
+    _, g2 = Optimisers.update(c2, s2, m, (α = ([0.1], nothing), γ = [1,10,100],))
+    @test only(m.α[1] .- g2.α[1]) ≈ 0.1
+    @test norm(m.γ .- g2.γ) ≈ 10
+    @test_throws DomainError Optimisers.update(c2, s2, m, (α = [0.1], γ = [1,10,NaN],))
+
+    c3 = ClipNorm(5, 1; throw=false)
+    _, g3 = Optimisers.update(c3, s2, m, (α = ([0.1], nothing), γ = [1,10,100],))
+    @test only(m.α[1] .- g3.α[1]) ≈ 0.1
+    @test norm(m.γ .- g3.γ, 1) ≈ 5
+    _, g3n = Optimisers.update(c3, s2, m, (α = nothing, γ = [1,10,Inf],))
+    @test isnan(g3n.γ[3])
+  end
+
+  @testset "Optimiser Updates" begin
+    opt = ADAM()
+    new_opt = ADAM(opt, eta = 9.f0)
+    @test new_opt.eta == 9.f0
   end
-  @test loss(rand(10, 10), w, w′) < 0.01
-end
 
-@testset "Optimiser Updates" begin
-  opt = ADAM()
-  new_opt = ADAM(opt, eta = 9.f0)
-  @test new_opt.eta == 9.f0
 end