Merge pull request #56 from SciML/loss

Implement l2 loss

Author: foldfelis

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -31,7 +31,7 @@ function train()
         Dense(64, 1),
     ) |> device
 
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+    loss(𝐱, 𝐲) = l₂loss(m(𝐱), 𝐲)
 
     opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
 

src/NeuralOperators.jl

@@ -14,6 +14,7 @@ module NeuralOperators
 
     include("Transform/Transform.jl")
     include("operator_kernel.jl")
+    include("loss.jl")
     include("model.jl")
     include("DeepONet.jl")
     include("subnets.jl")

src/loss.jl

@@ -0,0 +1,8 @@
+export l₂loss
+
+function l₂loss(𝐲̂, 𝐲; agg=mean, grid_normalize=true)
+    feature_dims = 2:(ndims(𝐲)-1)
+    loss = agg(.√(sum(abs2, 𝐲̂-𝐲, dims=feature_dims)))
+
+    return grid_normalize ? loss/prod(feature_dims) : loss
+end

test/loss.jl

@@ -0,0 +1,9 @@
+@testset "loss" begin
+    𝐲 = rand(1, 3, 3, 5)
+    𝐲̂ = rand(1, 3, 3, 5)
+
+    feature_dims = 2:3
+    loss = mean(.√(sum(abs2, 𝐲̂-𝐲, dims=feature_dims)))
+
+    @test l₂loss(𝐲̂, 𝐲) ≈ loss/prod(feature_dims)
+end

test/runtests.jl

@@ -4,19 +4,21 @@ using Flux
 using GeometricFlux
 using Graphs
 using Zygote
+using Statistics
 using Test
 
 CUDA.allowscalar(false)
 
 cuda_tests = [
-    "cuda",
+    "cuda.jl",
 ]
 
 tests = [
-    "Transform/Transform",
-    "operator_kernel",
-    "model",
-    "deeponet",
+    "Transform/Transform.jl",
+    "operator_kernel.jl",
+    "loss.jl",
+    "model.jl",
+    "deeponet.jl",
 ]
 
 if CUDA.functional()
@@ -27,7 +29,7 @@ end
 
 @testset "NeuralOperators.jl" begin
     for t in tests
-        include("$(t).jl")
+        include(t)
     end
 end