Replace ADAM with Adam

yuehhua · Jul 7, 2022 · f4db7ea · f4db7ea
1 parent 1221dc9
commit f4db7ea
Show file tree

Hide file tree

Showing 9 changed files with 21 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -76,7 +76,7 @@ And then train as a Flux model.
 
 ```julia
 loss(𝐱, 𝐲) = l₂loss(model(𝐱), 𝐲)
-opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+opt = Flux.Optimiser(WeightDecay(1f-4), Flux.Adam(1f-3))
 Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)
 ```
 
@@ -102,7 +102,7 @@ loss(xtrain, ytrain, sensor) = Flux.Losses.mse(model(xtrain, sensor), ytrain)
 evalcb() = @show(loss(xval, yval, grid))
 
 learning_rate = 0.001
-opt = ADAM(learning_rate)
+opt = Adam(learning_rate)
 parameters = params(model)
 Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb=evalcb)
 ```

diff --git a/docs/src/index.md b/docs/src/index.md
@@ -53,7 +53,7 @@ And then train as a Flux model.
 
 ```julia
 loss(𝐱, 𝐲) = l₂loss(model(𝐱), 𝐲)
-opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
+opt = Flux.Optimiser(WeightDecay(1f-4), Flux.Adam(1f-3))
 Flux.@epochs 50 Flux.train!(loss, params(model), data, opt)
 ```
 
@@ -80,7 +80,7 @@ loss(xtrain, ytrain, sensor) = Flux.Losses.mse(model(xtrain, sensor), ytrain)
 evalcb() = @show(loss(xval, yval, grid))
 
 learning_rate = 0.001
-opt = ADAM(learning_rate)
+opt = Adam(learning_rate)
 parameters = params(model)
 Flux.@epochs 400 Flux.train!(loss, parameters, [(xtrain, ytrain, grid)], opt, cb=evalcb)
 ```

diff --git a/example/Burgers/src/Burgers.jl b/example/Burgers/src/Burgers.jl
@@ -56,7 +56,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 500)
     model = FourierNeuralOperator(ch = (2, 64, 64, 64, 64, 64, 128, 1), modes = (16,),
                                   σ = gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
     learner = Learner(model, data, optimiser, loss_func,
@@ -88,7 +88,7 @@ function train_nomad(; n = 300, cuda = true, learning_rate = 0.001, epochs = 400
     grid = rand(collect(0:0.001:1), (280, 1024)) |> device
     gridval = rand(collect(0:0.001:1), (20, 1024)) |> device
 
-    opt = ADAM(learning_rate)
+    opt = Adam(learning_rate)
 
     m = NOMAD((1024, 1024), (2048, 1024), gelu, gelu) |> device
 

diff --git a/example/Burgers/src/Burgers_deeponet.jl b/example/Burgers/src/Burgers_deeponet.jl
@@ -26,7 +26,7 @@ function train_don(; n = 300, cuda = true, learning_rate = 0.001, epochs = 400)
 
     grid = collect(range(0, 1, length = 1024)') |> device
 
-    opt = ADAM(learning_rate)
+    opt = Adam(learning_rate)
 
     m = DeepONet((1024, 1024, 1024), (1, 1024, 1024), gelu, gelu) |> device
 

diff --git a/example/DoublePendulum/src/DoublePendulum.jl b/example/DoublePendulum/src/DoublePendulum.jl
@@ -93,7 +93,7 @@ function train(; cuda = true, Δt = 1, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 20)
     model = FourierNeuralOperator(ch = (2, 64, 64, 64, 64, 64, 128, 2), modes = (4, 16),
                                   σ = gelu)
     data = get_dataloader(Δt = Δt)
-    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
     learner = Learner(model, data, optimiser, loss_func,

diff --git a/example/FlowOverCircle/src/FlowOverCircle.jl b/example/FlowOverCircle/src/FlowOverCircle.jl
@@ -62,7 +62,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     model = MarkovNeuralOperator(ch = (1, 64, 64, 64, 64, 64, 1), modes = (24, 24),
                                  σ = gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
     learner = Learner(model, data, optimiser, loss_func,
@@ -92,7 +92,7 @@ function train_gno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
                   WithGraph(featured_graph, GraphKernel(Dense(2 * 16, 16, gelu), 16)),
                   Dense(16, 1))
     data = get_dataloader(batchsize = 16, flatten = true)
-    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
     learner = Learner(model, data, optimiser, loss_func,

diff --git a/example/SuperResolution/src/SuperResolution.jl b/example/SuperResolution/src/SuperResolution.jl
@@ -87,7 +87,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     model = MarkovNeuralOperator(ch = (1, 64, 64, 64, 64, 64, 1), modes = (24, 24),
                                  σ = gelu)
     data = get_dataloader()
-    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η₀))
+    optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
     learner = Learner(model, data, optimiser, loss_func,

diff --git a/test/model.jl b/test/model.jl
@@ -6,7 +6,7 @@
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
     data = [(𝐱, 𝐲)]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "MarkovNeuralOperator" begin
@@ -17,5 +17,5 @@ end
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
     data = [(𝐱, 𝐲)]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
diff --git a/test/operator_kernel.jl b/test/operator_kernel.jl
@@ -13,7 +13,7 @@
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 128, 1024, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "permuted 1D OperatorConv" begin
@@ -32,7 +32,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 1024, 128, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "1D OperatorKernel" begin
@@ -49,7 +49,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 128, 1024, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "permuted 1D OperatorKernel" begin
@@ -67,7 +67,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 1024, 128, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "2D OperatorConv" begin
@@ -83,7 +83,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 64, 22, 22, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "permuted 2D OperatorConv" begin
@@ -100,7 +100,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 22, 22, 64, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "2D OperatorKernel" begin
@@ -115,7 +115,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 64, 22, 22, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "permuted 2D OperatorKernel" begin
@@ -131,7 +131,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 22, 22, 64, 5))]
-    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.Adam())
 end
 
 @testset "SpectralConv" begin