Repository: FluxML/Optimisers.jl
Pull request #7 (2021-01-05 11:21:47)

Add some more Optimisers
- [x] Added Momentum
- [x] Added Nesterov
- [x] RMSProp


Patch:
diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
new file mode 100644
index 0000000..ce8d353
--- /dev/null
+++ b/.github/workflows/CompatHelper.yml
@@ -0,0 +1,16 @@
+name: CompatHelper
+
+on:
+  schedule:
+    - cron: '00 00 * * *'
+
+jobs:
+  CompatHelper:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Pkg.add("CompatHelper")
+        run: julia -e 'using Pkg; Pkg.add("CompatHelper")'
+      - name: CompatHelper.main()
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: julia -e 'using CompatHelper; CompatHelper.main()'
diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml
new file mode 100644
index 0000000..a0f3109
--- /dev/null
+++ b/.github/workflows/TagBot.yml
@@ -0,0 +1,16 @@
+name: TagBot
+on:
+  issue_comment:  # THIS BIT IS NEW
+    types:
+      - created
+  workflow_dispatch:
+jobs:
+  TagBot:
+    # THIS 'if' LINE IS NEW
+    if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'
+    # NOTHING BELOW HAS CHANGED
+    runs-on: ubuntu-latest
+    steps:
+      - uses: JuliaRegistries/TagBot@v1
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}
diff --git a/Manifest.toml b/Manifest.toml
index 74aed06..23989dc 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -3,30 +3,26 @@
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
-[[Distributed]]
-deps = ["Random", "Serialization", "Sockets"]
-uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
-
 [[Functors]]
 deps = ["MacroTools"]
-git-tree-sha1 = "f70dabffddbfe9b22e5f61800c7ca813f1e685df"
+git-tree-sha1 = "bd96c1974893293ecc8b7e9016954d8cbd840ee4"
 repo-rev = "master"
 repo-url = "https://github.com/FluxML/Functors.jl"
 uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 version = "0.1.0"
 
-[[InteractiveUtils]]
-deps = ["Markdown"]
-uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+[[Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 
-[[Logging]]
-uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+[[LinearAlgebra]]
+deps = ["Libdl"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[MacroTools]]
 deps = ["Markdown", "Random"]
-git-tree-sha1 = "f7d2e3f654af75f01ec49be82c231c382214223a"
+git-tree-sha1 = "6a8a2a625ab0dea913aba95c11370589e0239ff0"
 uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
-version = "0.5.5"
+version = "0.5.6"
 
 [[Markdown]]
 deps = ["Base64"]
@@ -37,17 +33,25 @@ deps = ["Serialization"]
 uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [[Requires]]
-deps = ["Test"]
-git-tree-sha1 = "f6fbf4ba64d295e146e49e021207993b6b48c7d1"
+deps = ["UUIDs"]
+git-tree-sha1 = "cfbac6c1ed70c002ec6361e7fd334f02820d6419"
 uuid = "ae029012-a4dd-5104-9daa-d747884805df"
-version = "0.5.2"
+version = "1.1.2"
+
+[[SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
-[[Sockets]]
-uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+[[SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
-[[Test]]
-deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
-uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+[[UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
diff --git a/Project.toml b/Project.toml
index 3c69d0d..62b70e6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,10 +5,17 @@ version = "0.1.0"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[compat]
+Functors = "0"
+Requires = "0.5, 1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Zygote"]
diff --git a/src/Optimisers.jl b/src/Optimisers.jl
index f0c0ae6..d72f235 100644
--- a/src/Optimisers.jl
+++ b/src/Optimisers.jl
@@ -5,4 +5,7 @@ using Functors: functor, fmap, isleaf
 include("interface.jl")
 include("rules.jl")
 
+export Descent, Momentum, Nesterov, RMSProp,
+       ADAM
+
 end # module
diff --git a/src/interface.jl b/src/interface.jl
index bd202c0..f04727d 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -11,12 +11,12 @@ function state(o, x)
   end
 end
 
-function _update(o, x, x̄, state)
-  x̄, state = apply(o, x, x̄, state)
-  return patch(x, x̄), state
+function _update(o, x, x̄, st)
+  x̄, st = apply(o, x, x̄, st)
+  return patch(x, x̄), st
 end
 
-function update(o, x, x̄, state)
+function update(o, x::T, x̄, state) where T
   if x̄ === nothing
     return x, state
   elseif isleaf(x)
diff --git a/src/rules.jl b/src/rules.jl
index ac7b1b2..f274bda 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -23,6 +23,69 @@ function (o::Descent)(m, m̄, st)
   update(o, m, m̄, st)
 end
 
+struct Momentum{T,S}
+  eta::T
+  rho::S
+end
+
+Momentum(η = 0.01, ρ = 0.9) = Momentum{typeof(η), typeof(ρ)}(η, ρ)
+
+function apply(o::Momentum, x, Δ, st)
+  η, ρ = o.eta, o.rho
+  v = st
+  v = @. ρ * v - η * Δ
+  Δ = @. -v
+  Δ, v
+end
+
+function (o::Momentum)(m, m̄, state)
+  update(o, m, m̄, state)
+end
+
+init(o::Momentum, x::AbstractArray) = zero(x)
+
+struct Nesterov{T,S}
+  eta::T
+  rho::S
+end
+
+Nesterov(η = 0.001, ρ = 0.9) = Nesterov{typeof(η), typeof(ρ)}(η, ρ)
+
+init(o::Nesterov, x::AbstractArray) = zero(x)
+
+function (o::Nesterov)(m, m̄, state)
+  update(o, m, m̄, state)
+end
+
+function apply(o::Nesterov, x, Δ, st)
+  η, ρ = o.eta, o.rho
+  v = st
+  d = @. ρ^2 * v - (1+ρ) * η * Δ
+  v = @. ρ*v - η*Δ
+  Δ = -d
+  Δ, v
+end
+
+struct RMSProp{T,S}
+  eta::T
+  rho::S
+end
+
+init(o::RMSProp, x::AbstractArray) = zero(x)
+RMSProp(η = 0.001, ρ = 0.9) = RMSProp{typeof(η), typeof(ρ)}(η, ρ)
+
+function apply(o::RMSProp, x, Δ, st)
+  η, ρ = o.eta, o.rho
+  acc = st
+  acc = ρ .* acc .+ (1 .- ρ) .* Δ.^2
+  Δ = Δ .* (η ./ (.√acc .+ ϵ))
+  Δ, acc
+end
+
+function (o::RMSProp)(m, m̄, state)
+  update(o, m, m̄, state)
+end
+
 struct ADAM{T,K}
   eta::T
   beta::Tuple{K,K}
@@ -45,3 +108,4 @@ function apply(o::ADAM, x, Δ, st)
   Δ =  mt ./ (1 .- βp[1]) ./ (.√(vt ./ (1f0 .- βp[2])) .+ ϵ) .* η
   return Δ, (mt, vt, βp .* β)
 end
+


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 7b8d313..c01ed9e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,21 @@
 using Optimisers, Test
+using Zygote, Random
+using Statistics
 
 @testset "Optimisers" begin
+  Random.seed!(84)
+  w′ = rand(3,3)
+  @testset for o in (Descent(0.1), Momentum(), Nesterov(), RMSProp(),
+                     ADAM(0.001, (0.9, 0.99)))
+    w = rand(3,3)
+    st = Optimisers.init(o,w)
+    loss(x, y) = mean((x .- y) .^ 2)
+    l = loss(w, w′)
+    for i = 1:10^4
+      gs = gradient(x -> loss(x,w′), w)
+      w, st = o(w, gs..., st)
+    end
+    @test loss(w, w′) < 0.01
+  end
 
 end