refactor example

Author: foldfelis

example/burgers.jl

@@ -13,25 +13,23 @@ end
 m = FourierNeuralOperator() |> device
 loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
 
+n_train = 1800
+n_test = 200
+batchsize = 100
 𝐱, 𝐲 = get_burgers_data(n=2048)
 
-n_train = 1000
 𝐱_train, 𝐲_train = 𝐱[:, :, 1:n_train], 𝐲[:, 1:n_train]
-loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=20, shuffle=true)
+loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)
 
-n_test = 100
 𝐱_test, 𝐲_test = 𝐱[:, :, end-n_test+1:end], 𝐲[:, end-n_test+1:end]
-loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=20, shuffle=false)
+loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)
 
-function loss_test()
-    l = 0f0
-    for (𝐱, 𝐲) in loader_test
-        𝐱, 𝐲 = device(𝐱), device(𝐲)
-        l += loss(𝐱, 𝐲)
-    end
-    @info "loss: $(l/length(loader_test))"
+function validate()
+    validation_losses = [loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test]
+    @info "loss: $(sum(validation_losses)/length(loader_test))"
 end
 
 data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
 opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
-Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=Flux.throttle(loss_test, 5)))
+call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
+Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))