separate get_dataloader

Author: yuehhua

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -3,7 +3,7 @@ module FlowOverCircle
 using WaterLily, LinearAlgebra, ProgressMeter, MLUtils
 using NeuralOperators, Flux, GeometricFlux, Graphs
 using CUDA, FluxTraining, BSON
-using GeometricFlux.GraphSignals: generate_coordinates
+using GeometricFlux.GraphSignals: generate_grid
 
 function circle(n, m; Re = 250) # copy from [WaterLily](https://github.com/weymouth/WaterLily.jl)
     # Set physical parameters
@@ -32,29 +32,21 @@ function gen_data(ts::AbstractRange)
     return 𝐩s
 end
 
-function get_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
-                        ratio::Float64 = 0.95, batchsize = 100, flatten = false)
+function get_mno_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
+                            ratio::Float64 = 0.95, batchsize = 100)
     data = gen_data(ts)
     𝐱, 𝐲 = data[:, :, :, 1:(end - 1)], data[:, :, :, 2:end]
     n = length(ts) - 1
-    grid = generate_coordinates(𝐱[1, :, :, 1])
-    grid = repeat(grid, outer = (1, 1, 1, n))
-    x_with_grid = vcat(𝐱, grid)
-
-    if flatten
-        x_with_grid = reshape(x_with_grid, size(x_with_grid, 1), :, n)
-        𝐲 = reshape(𝐲, 1, :, n)
-    end
 
-    data_train, data_test = splitobs(shuffleobs((x_with_grid, 𝐲)), at = ratio)
+    data_train, data_test = splitobs(shuffleobs((𝐱, 𝐲)), at = ratio)
 
     loader_train = DataLoader(data_train, batchsize = batchsize, shuffle = true)
     loader_test = DataLoader(data_test, batchsize = batchsize, shuffle = false)
 
     return loader_train, loader_test
 end
 
-function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
+function train_mno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     if cuda && CUDA.has_cuda()
         device = gpu
         CUDA.allowscalar(false)
@@ -66,7 +58,7 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
 
     model = MarkovNeuralOperator(ch = (1, 64, 64, 64, 64, 64, 1), modes = (24, 24),
                                  σ = gelu)
-    data = get_dataloader()
+    data = get_mno_dataloader()
     optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
 
@@ -79,6 +71,61 @@ function train(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     return learner
 end
 
+function batch_featured_graph(data, graph, batchsize)
+    tot_len = size(data)[end]
+    bch_data = FeaturedGraph[]
+    for i in 1:batchsize:tot_len
+        bch_rng = (i + batchsize >= tot_len) ? (i:tot_len) : (i:(i + batchsize - 1))
+        fg = FeaturedGraph(graph, nf = data[:, :, bch_rng], pf = data[:, :, bch_rng])
+        push!(bch_data, fg)
+    end
+
+    return bch_data
+end
+
+function batch_data(data, batchsize)
+    tot_len = size(data)[end]
+    bch_data = Array{Float32, 3}[]
+    for i in 1:batchsize:tot_len
+        bch_rng = (i + batchsize >= tot_len) ? (i:tot_len) : (i:(i + batchsize - 1))
+        push!(bch_data, data[:, :, bch_rng])
+    end
+
+    return bch_data
+end
+
+function get_gno_dataloader(; ts::AbstractRange = LinRange(100, 11000, 10000),
+                            ratio::Float64 = 0.95, batchsize = 8)
+    data = gen_data(ts)
+    𝐱, 𝐲 = data[:, :, :, 1:(end - 1)], data[:, :, :, 2:end]
+    n = length(ts) - 1
+
+    # generate graph
+    graph = Graphs.grid(size(data)[2:3])
+
+    # add grid coordinates
+    grid = generate_coordinates(𝐱[1, :, :, 1])
+    grid = repeat(grid, outer = (1, 1, 1, n))
+    𝐱 = vcat(𝐱, grid)
+
+    # flatten
+    𝐱, 𝐲 = reshape(𝐱, size(𝐱, 1), :, n), reshape(𝐲, 1, :, n)
+
+    data_train, data_test = splitobs(shuffleobs((𝐱, 𝐲)), at = ratio)
+
+    batched_train_X = batch_featured_graph(data_train[1], graph, batchsize)
+    batched_test_X = batch_featured_graph(data_test[1], graph, batchsize)
+    batched_train_y = batch_data(data_train[2], batchsize)
+    batched_test_y = batch_data(data_test[2], batchsize)
+
+    loader_train = DataLoader((batched_train_X, batched_train_y), batchsize = -1,
+                              shuffle = true)
+    loader_test = DataLoader((batched_test_X, batched_test_y), batchsize = -1,
+                             shuffle = false)
+
+    return loader_train, loader_test
+end
+
 function train_gno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
     if cuda && CUDA.has_cuda()
         device = gpu
@@ -91,23 +138,17 @@ function train_gno(; cuda = true, η₀ = 1.0f-3, λ = 1.0f-4, epochs = 50)
 
     grid_dim = 2
     edge_dim = 2(grid_dim + 1)
-    featured_graph = FeaturedGraph(grid([96, 64]))
-    model = Chain(Flux.SkipConnection(Dense(grid_dim + 1, 16), vcat),
-                 # size(x) = (19, 6144, 8)
-                  WithGraph(featured_graph,
-                            GraphKernel(Dense(edge_dim, abs2(16), gelu), 16)),
-                  WithGraph(featured_graph,
-                            GraphKernel(Dense(edge_dim, abs2(16), gelu), 16)),
-                  WithGraph(featured_graph,
-                            GraphKernel(Dense(edge_dim, abs2(16), gelu), 16)),
-                  WithGraph(featured_graph,
-                            GraphKernel(Dense(edge_dim, abs2(16), gelu), 16)),
-                  x -> x[1:end-3, :, :],
+    model = Chain(GraphParallel(node_layer = Dense(grid_dim + 1, 16)),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  GraphKernel(Dense(edge_dim, abs2(16), gelu), 16),
+                  node_feature,
                   Dense(16, 1))
 
     optimiser = Flux.Optimiser(WeightDecay(λ), Flux.Adam(η₀))
     loss_func = l₂loss
-    data = get_dataloader(batchsize = 8, flatten = true)
+    data = get_gno_dataloader()
     learner = Learner(model, data, optimiser, loss_func,
                       ToDevice(device, device),
                       Checkpointer(joinpath(@__DIR__, "../model/")))

src/graph_kernel.jl

@@ -44,23 +44,25 @@ end
 
 function (l::GraphKernel)(el::NamedTuple, X::AbstractArray, E::AbstractArray)
     GraphSignals.check_num_nodes(el.N, X)
-    # GraphSignals.check_num_edges(el.E, E)
+    GraphSignals.check_num_edges(el.E, E)
     _, V, _ = GeometricFlux.propagate(l, el, E, X, nothing, mean, nothing, nothing)
 
     return V
 end
 
-# (wg::WithGraph{<:GraphKernel})(Vt::AbstractArray, E::AbstractArray) = wg(nothing, Vt, E), E
-
-function (wg::WithGraph{<:GraphKernel})(input::AbstractArray)
-    el = wg.graph
-    Vt, X_with_grid = input[1:end-3, :, :], input[end-2:end, :, :]
+# For variable graph
+function (l::GraphKernel)(fg::AbstractFeaturedGraph)
+    nf = node_feature(fg)
+    pf = positional_feature(fg)
+    GraphSignals.check_num_nodes(fg, nf)
+    GraphSignals.check_num_nodes(fg, pf)
+    el = GeometricFlux.GraphSignals.to_namedtuple(fg)
 
     # node features + positional features as edge features
-    E = vcat(GeometricFlux._gather(X_with_grid, el.xs),
-             GeometricFlux._gather(X_with_grid, el.nbrs))
+    ef = vcat(GeometricFlux._gather(pf, el.xs), GeometricFlux._gather(pf, el.nbrs))
+    _, V, _ = GeometricFlux.propagate(l, el, ef, nf, nothing, mean, nothing, nothing)
 
-    return vcat(wg.layer(el, Vt, E), X_with_grid)
+    return ConcreteFeaturedGraph(fg, nf = V)
 end
 
 function Base.show(io::IO, l::GraphKernel)

test/graph_kernel.jl

@@ -9,40 +9,16 @@
     κ = Dense(2(coord_dim + channel), abs2(channel), relu)
     κ_in_dim, κ_out_dim = 2(coord_dim + channel), abs2(channel)
 
-    @testset "pass edge features" begin
-        E = rand(Float32, 2(coord_dim + channel), ne(graph), batch_size)
-        l = WithGraph(FeaturedGraph(graph), GraphKernel(κ, channel))
-        @test repr(l.layer) ==
-              "GraphKernel(Dense($κ_in_dim => $κ_out_dim, relu), channel=$channel)"
-        @test size(l(𝐱, E)) == (channel, nv(graph), batch_size)
-
-        g = Zygote.gradient(() -> sum(l(𝐱, E)), Flux.params(l))
-        @test length(g.grads) == 3
-    end
-
-    @testset "pass positional features" begin
-        pf = rand(Float32, coord_dim, nv(graph), batch_size)
-        pf = vcat(𝐱, pf)
-        fg = FeaturedGraph(graph)
-        l = WithGraph(fg, GraphKernel(κ, channel))
-        @test repr(l.layer) ==
-              "GraphKernel(Dense($κ_in_dim => $κ_out_dim, relu), channel=$channel)"
-        @test size(l(pf, 𝐱, nothing)) == (channel, nv(graph), batch_size)
-
-        g = Zygote.gradient(() -> sum(l(pf, 𝐱, nothing)), Flux.params(l))
-        @test length(g.grads) == 4
-    end
-
-    @testset "pass positional features by FeaturedGraph" begin
+    @testset "layer without graph" begin
         pf = rand(Float32, coord_dim, nv(graph), batch_size)
         pf = vcat(𝐱, pf)
-        fg = FeaturedGraph(graph, pf = pf)
-        l = WithGraph(fg, GraphKernel(κ, channel))
-        @test repr(l.layer) ==
-              "GraphKernel(Dense($κ_in_dim => $κ_out_dim, relu), channel=$channel)"
-        @test size(l(𝐱)) == (channel, nv(graph), batch_size)
+        l = GraphKernel(κ, channel)
+        fg = FeaturedGraph(graph, nf = 𝐱, pf = pf)
+        fg_ = l(fg)
+        @test size(node_feature(fg_)) == (channel, nv(graph), batch_size)
+        @test_throws MethodError l(𝐱)
 
-        g = Zygote.gradient(() -> sum(l(𝐱)), Flux.params(l))
-        @test length(g.grads) == 3
+        g = gradient(() -> sum(node_feature(l(fg))), Flux.params(l))
+        @test length(g.grads) == 5
     end
 end