feat: setup basic flux code

[skip ci] [skip docs]

bench/comparison.md

@@ -1,41 +1,49 @@
-## NeuralOperators.jl (Lux)
+# NeuralOperators.jl Benchmarks
 
-## FNO
+## Fourier Neural Operators
 
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-| 1 | 14.173699999999998 ms | 755.1466 ms |
-| 2 | 29.118399999999998 ms | 1407.2298 ms |
-| 3 | 37.6924 ms | 2367.5004999999996 ms |
-| 4 | 41.431400000000004 ms | 3035.1971 ms |
-| 5 | 59.305 ms | 3456.1902999999998 ms |
+### Lux.jl (Julia)
 
-## FNO (python: neuraloperator)
+| #layers | Forward | Train: 10 epochs |
+|:------- |:--------------------- |:--------------------- |
+| 1 | 14.173699999999998 ms | 755.1466 ms |
+| 2 | 29.118399999999998 ms | 1407.2298 ms |
+| 3 | 37.6924 ms | 2367.5004999999996 ms |
+| 4 | 41.431400000000004 ms | 3035.1971 ms |
+| 5 | 59.305 ms | 3456.1902999999998 ms |
 
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-| 1 | 5.731542900000932 ms | 17.667421199992532 ms |
-| 2 | 7.833489999989979 ms | 25.585920999990776 ms |
-| 3 | 10.18306370000937 ms | 33.69801080002799 ms |
-| 4 | 12.33892210002523 ms | 41.98180860001594 ms |
-| 5 | 14.732645300013246 ms | 50.13744520000182 ms |
+### Flux.jl (Julia)
+
+### neuraloperator (Python)
+
+| #layers | Forward | Train: 10 epochs |
+|:------- |:---------------------- |:---------------------- |
+| 1 | 5.731542900000932 ms | 17.667421199992532 ms |
+| 2 | 7.833489999989979 ms | 25.585920999990776 ms |
+| 3 | 10.18306370000937 ms | 33.69801080002799 ms |
+| 4 | 12.33892210002523 ms | 41.98180860001594 ms |
+| 5 | 14.732645300013246 ms | 50.13744520000182 ms |
 
 ## DeepONet
 
-| #layers | Forward | Train: 10 epochs | 
-| --- | --- | --- | 
-| 1 | 3.3952750000000003 ms | 76.604576 ms | 
-| 2 | 4.360458 ms | 104.460251 ms | 
-| 3 | 5.6310780000000005 ms | 149.148633 ms | 
-| 4 | 7.199777 ms | 178.464657 ms | 
-| 5 | 7.8226819999999995 ms | 193.760173 ms |
-
-## DeepONet (python: deepxde)
-
-| #layers | Forward | Train: 10 epochs |
-| --- | --- | --- |
-| 1 | 0.7689221948385239 ms | 25.76469287276268 ms |
-| 2 | 0.7733150571584702 ms | 32.17746138572693 ms |
-| 3 | 0.8474267274141312 ms | 36.93301998078823 ms |
-| 4 | 1.0069304704666138 ms | 45.45578710734844 ms |
-| 5 | 1.406572386622429 ms | 59.06449243426323 ms |
+### Lux.jl (Julia)
+
+| #layers | Forward | Train: 10 epochs |
+|:------- |:--------------------- |:---------------- |
+| 1 | 3.3952750000000003 ms | 76.604576 ms |
+| 2 | 4.360458 ms | 104.460251 ms |
+| 3 | 5.6310780000000005 ms | 149.148633 ms |
+| 4 | 7.199777 ms | 178.464657 ms |
+| 5 | 7.8226819999999995 ms | 193.760173 ms |
+
+### Flux.jl (Julia)
+
+### deepxde (Python)
+
+| #layers | Forward | Train: 10 epochs |
+|:------- |:---------------------- |:--------------------- |
+| 1 | 0.7689221948385239 ms | 25.76469287276268 ms |
+| 2 | 0.7733150571584702 ms | 32.17746138572693 ms |
+| 3 | 0.8474267274141312 ms | 36.93301998078823 ms |
+| 4 | 1.0069304704666138 ms | 45.45578710734844 ms |
+| 5 | 1.406572386622429 ms | 59.06449243426323 ms |

bench/flux.jl

@@ -0,0 +1,66 @@
+using ThreadPinning
+pinthreads(:cores)
+threadinfo()
+
+using BenchmarkTools, NeuralOperators, Random, Optimisers, Zygote
+
+# TODO: Add training code
+
+# FNO
+n_points = 128
+batch_size = 64
+
+x = rand(Float32, 1, n_points, batch_size);
+y = rand(Float32, 1, n_points, batch_size);
+data = [(x, y)];
+t_fwd = zeros(5)
+t_train = zeros(5)
+
+for i in 1:5
+ chs = (1, 128, fill(64, i)..., 128, 1)
+ model = FourierNeuralOperator(; ch=chs, modes=(16,), σ=gelu)
+ # model(x) # TTFX
+
+ # t_fwd[i] = @belapsed $model($x)
+
+ # t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
+end
+
+println("\n## FNO (Flux NeuralOperators.jl)")
+print("| #layers | Forward | Train: 10 epochs | \n")
+print("| --- | --- | --- | \n")
+for i in 1:5
+ print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
+end
+
+# DeepONets
+eval_points = 128
+batch_size = 64
+dim_y = 1
+m = 32
+
+u = rand(Float32, m, batch_size);
+y = rand(Float32, dim_y, eval_points, batch_size);
+
+g = rand(Float32, eval_points, batch_size);
+
+data = [((u, y), g)]
+t_fwd = zeros(5)
+t_train = zeros(5)
+for i in 1:5
+ ch_branch = (m, fill(64, i)..., 128)
+ ch_trunk = (dim_y, fill(64, i)..., 128)
+ model = DeepONet(ch_branch, ch_trunk)
+ # model(u, y) # TTFX
+
+ # t_fwd[i] = @belapsed $model($u, $y)
+
+ # t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
+end
+
+println("\n## DeepONet (Flux NeuralOperators.jl)")
+print("| #layers | Forward | Train: 10 epochs | \n")
+print("| --- | --- | --- | \n")
+for i in 1:5
+ print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
+end

bench/flux/Project.toml

@@ -0,0 +1,5 @@
+[deps]
+NeuralOperators = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+ThreadPinning = "811555cd-349b-4f26-b7bc-1f208b848042"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

bench/lux.jl

@@ -34,14 +34,14 @@ for i in 1:5
  chs = (1, 128, fill(64, i)..., 128, 1)
  model = FourierNeuralOperator(gelu; chs=chs, modes=(16,))
  ps, st = Lux.setup(rng, model)
- _ = model(x, ps, st) # TTFX
+ model(x, ps, st) # TTFX
 
  t_fwd[i] = @belapsed $model($x, $ps, $st)
 
  t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
 end
 
-println("\n ### FNO")
+println("\n## FNO")
 print("| #layers | Forward | Train: 10 epochs | \n")
 print("| --- | --- | --- | \n")
 for i in 1:5
@@ -67,18 +67,16 @@ for i in 1:5
  ch_trunk = (dim_y, fill(64, i)..., 128)
  model = DeepONet(; branch=ch_branch, trunk=ch_trunk)
  ps, st = Lux.setup(rng, model)
- _ = model((u, y), ps, st) # TTFX
+ model((u, y), ps, st) # TTFX
 
  t_fwd[i] = @belapsed $model(($u, $y), $ps, $st)
 
  t_train[i] = @belapsed train!($model, $ps, $st, $data; epochs=10)
 end
 
-println("\n ### DeepONet")
+println("\n## DeepONet")
 print("| #layers | Forward | Train: 10 epochs | \n")
 print("| --- | --- | --- | \n")
 for i in 1:5
  print("| $i | $(t_fwd[i] * 1000) ms | $(t_train[i] * 1000) ms | \n")
 end
-
-