add GNO example

fix

Author: yuehhua

example/SuperResolution/Project.toml

@@ -4,11 +4,14 @@ uuid = "a8258e1f-331c-4af2-83e9-878628278453"
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+GeometricFlux = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 WaterLily = "ed894a53-35f9-47f1-b17f-85db9237eebd"
 

example/SuperResolution/src/SuperResolution.jl

@@ -4,8 +4,11 @@ using NeuralOperators
 using Flux
 using Flux.Losses: mse
 using Flux.Data: DataLoader
+using GeometricFlux
+using Graphs
 using CUDA
 using JLD2
+using ProgressMeter: Progress, next!
 
 include("data.jl")
 include("models.jl")

example/SuperResolution/src/data.jl

@@ -15,20 +15,20 @@ function circle(n, m; Re=250)
     Simulation((n+2, m+2), [U, 0.], R; ν, body)
 end
 
-function gen_data(ts::AbstractRange)
+function gen_data(ts::AbstractRange, T=Float32)
     n, m = 2 * 3(2^5), 2 * 2^6
     circ = circle(n, m)
 
-    𝐩s = Array{Float32}(undef, 1, n, m, length(ts))
+    𝐩s = Array{T}(undef, 1, n, m, length(ts))
     for (i, t) in enumerate(ts)
         sim_step!(circ, t)
-        𝐩s[1, :, :, i] .= Float32.(circ.flow.p)[2:end-1, 2:end-1]
+        𝐩s[1, :, :, i] .= T.(circ.flow.p)[2:end-1, 2:end-1]
     end
 
     return 𝐩s
 end
 
-function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Float64=0.95, batchsize=100)
+function get_dataloader(; ts::AbstractRange=LinRange(100, 11000, 10000), ratio::Real=0.95, batchsize=100)
     data = gen_data(ts)
 
     n_train, n_test = floor(Int, length(ts)*ratio), floor(Int, length(ts)*(1-ratio))

example/SuperResolution/src/models.jl

@@ -26,7 +26,7 @@ function train_mno(; cuda=true, η=1f-3, λ=1f-4, epochs=50)
     opt = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η))
 
     # parameters
-    ps = Flux.params(model)
+    ps = Flux.params(m)
 
     # training
     min_loss = Inf32
@@ -37,10 +37,10 @@ function train_mno(; cuda=true, η=1f-3, λ=1f-4, epochs=50)
         progress = Progress(length(loader_train))
 
         for (𝐱, 𝐲) in loader_train
-            grad = gradient(() -> loss(model, 𝐱 |> device, 𝐲 |> device), ps)
+            grad = gradient(() -> loss(m, 𝐱 |> device, 𝐲 |> device), ps)
             Flux.Optimise.update!(opt, ps, grad)
-            train_loss = loss(model, loader_train, device)
-            test_loss = loss(model, loader_test, device)
+            train_loss = loss(m, loader_train, device)
+            test_loss = loss(m, loader_test, device)
 
             # progress meter
             next!(progress; showvalues=[
@@ -49,12 +49,78 @@ function train_mno(; cuda=true, η=1f-3, λ=1f-4, epochs=50)
             ])
 
             if test_loss ≤ min_loss
-                update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
+                update_model!(joinpath(@__DIR__, "../model/mno.jld2"), m)
+                min_loss = test_loss
             end
 
             train_steps += 1
         end
     end
 
     return m
-end
+end
+
+function train_gno(; channel=64, cuda=true, η=1f-3, λ=1f-4, epochs=50)
+    # GPU config
+    if cuda && CUDA.has_cuda()
+        device = gpu
+        CUDA.allowscalar(false)
+        @info "Training on GPU"
+    else
+        device = cpu
+        @info "Training on CPU"
+    end
+
+    @info "gen data... "
+    @time loader_train, loader_test = get_dataloader()
+
+    # build model
+    g = grid([12, 8])
+    fg = FeaturedGraph(g)
+
+    m = Chain(
+        Dense(1, 64),
+        WithGraph(fg, GraphKernel(Dense(2channel, channel, gelu), channel)),
+        WithGraph(fg, GraphKernel(Dense(2channel, channel, gelu), channel)),
+        WithGraph(fg, GraphKernel(Dense(2channel, channel, gelu), channel)),
+        WithGraph(fg, GraphKernel(Dense(2channel, channel, gelu), channel)),
+        Dense(64, 1),
+    ) |> device
+
+    # optimizer
+    opt = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η))
+    
+    # parameters
+    ps = Flux.params(m)
+
+    # training
+    min_loss = Inf32
+    train_steps = 0
+    @info "Start Training, total $(epochs) epochs"
+    for epoch = 1:epochs
+        @info "Epoch $(epoch)"
+        progress = Progress(length(loader_train))
+
+        for (𝐱, 𝐲) in loader_train
+            grad = gradient(() -> loss(m, 𝐱 |> device, 𝐲 |> device), ps)
+            Flux.Optimise.update!(opt, ps, grad)
+            train_loss = loss(m, loader_train, device)
+            test_loss = loss(m, loader_test, device)
+
+            # progress meter
+            next!(progress; showvalues=[
+                (:train_loss, train_loss),
+                (:test_loss, test_loss)
+            ])
+
+            if test_loss ≤ min_loss
+                update_model!(joinpath(@__DIR__, "../model/gno.jld2"), m)
+                min_loss = test_loss
+            end
+
+            train_steps += 1
+        end
+    end
+
+    return m
+end