refactor SuperResolution example

Author: yuehhua

example/SuperResolution/src/SuperResolution.jl

@@ -2,56 +2,20 @@ module SuperResolution
 
 using NeuralOperators
 using Flux
+using Flux.Losses: mse
+using Flux.Data: DataLoader
 using CUDA
 using JLD2
 
 include("data.jl")
+include("models.jl")
 
 function update_model!(model_file_path, model)
     model = cpu(model)
     jldsave(model_file_path; model)
     @warn "model updated!"
 end
 
-function train()
-    if has_cuda()
-        @info "CUDA is on"
-        device = gpu
-        CUDA.allowscalar(false)
-    else
-        device = cpu
-    end
-
-    m = Chain(
-        Dense(1, 64),
-        OperatorKernel(64=>64, (24, 24), gelu),
-        OperatorKernel(64=>64, (24, 24), gelu),
-        OperatorKernel(64=>64, (24, 24), gelu),
-        OperatorKernel(64=>64, (24, 24), gelu),
-        Dense(64, 1),
-    ) |> device
-
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
-
-    opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
-
-    @info "gen data... "
-    @time loader_train, loader_test = get_dataloader()
-
-    losses = Float32[]
-    function validate()
-        validation_loss = sum(loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test)/length(loader_test)
-        @info "loss: $validation_loss"
-
-        push!(losses, validation_loss)
-        (losses[end] == minimum(losses)) && update_model!(joinpath(@__DIR__, "../model/model.jld2"), m)
-    end
-    call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
-
-    data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
-    Flux.@epochs 50 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))
-end
-
 function get_model()
     f = jldopen(joinpath(@__DIR__, "../model/model.jld2"))
     model = f["model"]
@@ -60,4 +24,7 @@ function get_model()
     return model
 end
 
+loss(m, 𝐱, 𝐲) = mse(m(𝐱), 𝐲)
+loss(m, loader::DataLoader, device) = sum(loss(m, 𝐱 |> device, 𝐲 |> device) for (𝐱, 𝐲) in loader)/length(loader)
+
 end

example/SuperResolution/src/models.jl

@@ -0,0 +1,60 @@
+function train_mno(; 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
+    m = Chain(
+        Dense(1, 64),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        OperatorKernel(64=>64, (24, 24), gelu),
+        Dense(64, 1),
+    ) |> device
+
+    # optimizer
+    opt = Flux.Optimiser(WeightDecay(λ), Flux.ADAM(η))
+    
+    # parameters
+    ps = Flux.params(model)
+
+    # 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(model, 𝐱 |> device, 𝐲 |> device), ps)
+            Flux.Optimise.update!(opt, ps, grad)
+            train_loss = loss(model, loader_train, device)
+            test_loss = loss(model, 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/model.jld2"), m)
+            end
+
+            train_steps += 1
+        end
+    end
+
+    return m
+end