Merge pull request #5 from una-auxme/dev

v0.2.0

Author: JulianTrommer

Project.toml

@@ -1,7 +1,7 @@
 name = "GraphNetCore"
 uuid = "7809f980-de1b-4f9a-8451-85f041491431"
 authors = ["JT <julian.trommer@uni-a.de>"]
-version = "0.1.1"
+version = "0.2.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -29,11 +29,11 @@ KernelAbstractions = "0.9"
 Lux = "0.5"
 LuxCUDA = "0.3.1"
 NNlib = "0.9"
+Random = "1"
+Statistics = "1"
 Tullio = "0.3.7"
 Zygote = "0.6"
 cuDNN = "1.1.1"
-Random = "1"
-Statistics = "1"
 julia = "1.9"
 
 [extras]

src/graph_network.jl

@@ -31,9 +31,9 @@ mutable struct GraphNetwork
     model
     ps
     st
-    e_norm::NormaliserOnline
+    e_norm::Union{NormaliserOffline, NormaliserOnline}
     n_norm::Dict{String, Union{NormaliserOffline, NormaliserOnline}}
-    o_norm::NormaliserOnline
+    o_norm::Dict{String, Union{NormaliserOffline, NormaliserOnline}}
 end
 
 function build_mlp(input_size::T, latent_size::T, output_size::T, hidden_layers::T; layer_norm=true, dev=cpu) where T <: Integer
@@ -198,22 +198,20 @@ Loads the [`GraphNetwork`](@ref) from the latest checkpoint at the given path.
 - `df_train`: [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl) DataFrame containing the train losses at the checkpoints.
 - `df_valid`: [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl) DataFrame containing the validation losses at the checkpoints (only improvements are saved).
 """
-function load(quantities, dims, norms, output, message_steps, ls, hl, opt, device::Function, path::String)
+function load(quantities, dims, e_norms, n_norms, o_norms, output, message_steps, ls, hl, opt, device::Function, path::String)
     if isfile(joinpath(path, "checkpoints"))
         step = parse(Int, readlines(joinpath(path, "checkpoints"))[end])
         ps_data, ps_axes, st, e_norm, n_norm, o_norm, opt_state, df_train, df_valid = load(joinpath(path, "checkpoint_$step.jld2"), "ps_data", "ps_axes", "st", "e_norm", "n_norm", "o_norm", "opt_state", "df_train", "df_valid")
 
         ps = ComponentArray(ps_data, ps_axes) |> device
         st = st |> device
 
-        en = NormaliserOnline(e_norm, device)
-        for (k, n) in n_norm
-            norms[k] = NormaliserOnline(n, device)
-        end
-        on = NormaliserOnline(o_norm, device)
+        en = deserialize(e_norm, device)
+        nn = deserialize(n_norm, device)
+        on = deserialize(o_norm, device)
 
         model = build_model(quantities, dims, output, message_steps, ls, hl, device)
-        gn = GraphNetwork(model, ps, st, en, norms, on)
+        gn = GraphNetwork(model, ps, st, en, nn, on)
 
         if !isnothing(opt)
             return gn, nothing, df_train, df_valid
@@ -227,7 +225,7 @@ function load(quantities, dims, norms, output, message_steps, ls, hl, opt, devic
         ps = ComponentArray(ps) |> device
         st = st |> device
 
-        gn = GraphNetwork(model, ps, st, NormaliserOnline(dims + 1, device), norms, NormaliserOnline(output, device))
+        gn = GraphNetwork(model, ps, st, e_norms, n_norms, o_norms)
 
         return gn, nothing, DataFrame(step=Integer[], loss=Float32[]), DataFrame(step=Integer[], loss=Float32[])
     end

src/normaliser.jl

@@ -26,7 +26,11 @@ function NormaliserOffline(data_min::Float32, data_max::Float32)
     NormaliserOffline(data_min, data_max, 0.0f0, 1.0f0)
 end
 
-(n::NormaliserOffline)(F, is_training = false) = minmaxnorm(F, n.data_min, n.data_max, n.target_min, n.target_max)
+function NormaliserOffline(d::Dict{String, Any})
+    NormaliserOffline(d["data_min"], d["data_max"], d["target_min"], d["target_max"])
+end
+
+(n::NormaliserOffline)(F) = minmaxnorm(F, n.data_min, n.data_max, n.target_min, n.target_max)
 
 """
     inverse_data(n::NormaliserOffline, data)
@@ -153,9 +157,7 @@ end
 function serialize(ns::Dict{String, Union{NormaliserOffline, NormaliserOnline}})
     result = Dict{String, Any}()
     for (k, n) in ns
-        if typeof(n) == NormaliserOnline
-            result[k] = serialize(n)
-        end
+        result[k] = serialize(n)
     end
     return result
 end
@@ -170,3 +172,24 @@ function serialize(n::NormaliserOnline)
         "acc_sum_squared" => cpu_device()(n.acc_sum_squared)
     )
 end
+
+function serialize(n::NormaliserOffline)
+    return Dict{String, Any}(
+        "data_min" => n.data_min,
+        "data_max" => n.data_max,
+        "target_min" => n.target_min,
+        "target_max" => n.target_max
+    )
+end
+
+function deserialize(n::Dict{String, Any}, device::Function)
+    if haskey(n, "max_accumulations")
+        return NormaliserOnline(n, device)
+    elseif haskey(n, "data_min")
+        return NormaliserOffline(n)
+    else
+        features = keys(n)
+        norms = deserialize.(values(n), device)
+        return Dict{String, Union{NormaliserOffline, NormaliserOnline}}(features .=> norms)
+    end
+end

src/utils.jl

@@ -104,7 +104,7 @@ Calculates the mean squared error of the given arguments with [Tullio](https://g
 - The calculated mean squared error.
 """
 mse_reduce(target, output) = begin
-    @assert ndims(target) == 2 && ndims(output) == 2 "Only supported dimension is 2: dims = (target => $(dims(target)), output => $(dims(output))"
+    @assert ndims(target) == 2 && ndims(output) == 2 "Only supported number of dimensions is 2: dims = (target => $(ndims(target)), output => $(ndims(output)))"
     @tullio R[x] := (target[y, x] - output[y, x]) ^ 2
 end