diff --git a/LICENSE b/LICENSE index 00631cf..7f1b524 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,7 @@ MIT License Copyright (c) 2023 Julian Trommer +Copyright (c) 2026 Josef Kircher Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/Project.toml b/Project.toml index bdcb0e7..0ddb213 100644 --- a/Project.toml +++ b/Project.toml @@ -1,45 +1,51 @@ name = "GraphNetCore" uuid = "7809f980-de1b-4f9a-8451-85f041491431" -authors = ["JT "] -version = "0.3.1" +authors = ["Julian Trommer "] +version = "0.4.0" [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] +Adapt = "4.2" Aqua = "0.8" -CUDA = "5" +CUDA = "5.9.6" ComponentArrays = "0.15" DataFrames = "1.6" -ForwardDiff = "0.10" -JLD2 = "0.4" +Enzyme = "0.13.73" +JLD2 = "0.6" KernelAbstractions = "0.9" -Lux = "0.5" -LuxCUDA = "0.3" +Lux = "1.13" NNlib = "0.9" Random = "1" +Reactant = "0.2.169" +Setfield = "1.1.2" Statistics = "1" -Tullio = "0.3.7" -Zygote = "0.6" -cuDNN = "1.3" Test = "1" -julia = "1.10" +Tullio = "0.3.7" +Zygote = "0.6, 0.7" +cuDNN = "1.4.5" +julia = "1.11" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] diff --git a/src/GraphNetCore.jl b/src/GraphNetCore.jl index a2a7737..f6b029a 100644 --- a/src/GraphNetCore.jl +++ b/src/GraphNetCore.jl @@ -5,8 +5,9 @@ module GraphNetCore -using CUDA -using Lux, LuxCUDA +using Lux +using CUDA, cuDNN +using Reactant using Tullio using Random @@ -23,7 +24,7 @@ export NormaliserOffline, NormaliserOfflineMinMax, NormaliserOfflineMeanStd, NormaliserOnline # graph_network.jl -export build_model, step!, save!, load +export build_model, step!, set_training!, save!, load # normaliser.jl export inverse_data # utils.jl diff --git a/src/feature_graph.jl b/src/feature_graph.jl index ffe8716..49606af 100644 --- a/src/feature_graph.jl +++ b/src/feature_graph.jl @@ -3,6 +3,8 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # +using NNlib + """ FeatureGraph(nf, ef, senders, receivers) @@ -14,33 +16,16 @@ Data structure that is used as an input for the [`GraphNetwork`](@ref). - `senders`: List of nodes in the mesh where graph edges start. - `receivers`: List of nodes in the mesh where graph edges end. """ -mutable struct FeatureGraph{F <: AbstractArray, T <: AbstractArray} - nf::F - ef::F +mutable struct FeatureGraph{T <: AbstractArray} + nf::Any + ef::Any senders::T receivers::T end -""" - update_features!(g; nf, ef) - -Updates the node and edge features of the given [`FeatureGraph`](@ref). - -## Arguments -- `g`: [`FeatureGraph`](@ref) that should be updated. - -## Keyword Arguments -- `nf`: Updated node features. -- `ef`: Updated edge features. - -## Returns -- Updated graph as a [`FeatureGraph`](@ref) struct. -""" -function update_features!(g::FeatureGraph; nf, ef) - g.nf = nf - g.ef = ef - - return g +function FeatureGraph(fg::FeatureGraph; nf = fg.nf, ef = fg.ef, + senders = fg.senders, receivers = fg.receivers) + return FeatureGraph(nf, ef, senders, receivers) end """ @@ -76,5 +61,6 @@ Aggregates the node features based on the given [`FeatureGraph`](@ref) and updat """ @inline function aggregate_node_features(graph::FeatureGraph, updated_edge_features) return vcat(graph.nf, - NNlib.scatter(+, updated_edge_features, graph.receivers; dstsize = size(graph.nf))) + NNlib.scatter(+, updated_edge_features, graph.receivers; + dstsize = size(graph.nf), init = zero(eltype(graph.nf)))) end diff --git a/src/graph_net_blocks.jl b/src/graph_net_blocks.jl index cbf3f05..9687620 100644 --- a/src/graph_net_blocks.jl +++ b/src/graph_net_blocks.jl @@ -3,91 +3,35 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # -struct Encoder{T <: NamedTuple, N <: Lux.NAME_TYPE} <: - Lux.AbstractExplicitContainerLayer{(:layers,)} - layers::T - name::N +struct Encoder{N, E} <: Lux.AbstractLuxContainerLayer{(:node_layer, :edge_layer)} + node_layer::N + edge_layer::E end -function Encoder(node_model, edge_model; name::Lux.NAME_TYPE = nothing) - fields = (Symbol("node_model_fn"), Symbol("edge_model_fn")) - - return Encoder(NamedTuple{fields}((node_model, edge_model)), name) -end - -function (e::Encoder)(graph::FeatureGraph, ps, st::NamedTuple{fields}) where {fields} - encode!(e.layers, graph, ps, st) -end - -function encode!( - layers::NamedTuple{fields}, graph, ps, st::NamedTuple{fields}) where {fields} - nf, stn = layers[:node_model_fn](graph.nf, ps[:node_model_fn], st[:node_model_fn]) - ef, ste = layers[:edge_model_fn](graph.ef, ps[:edge_model_fn], st[:edge_model_fn]) - new_st = NamedTuple{fields}((stn, ste)) - - return update_features!(graph; nf = nf, ef = ef), new_st -end - -struct Processor{T <: NamedTuple, N <: Lux.NAME_TYPE} <: - Lux.AbstractExplicitContainerLayer{(:layers,)} - layers::T - name::N -end - -function Processor(node_model, edge_model; name::Lux.NAME_TYPE = nothing) - fields = (Symbol("node_model_fn"), Symbol("edge_model_fn")) - - return Processor(NamedTuple{fields}((node_model, edge_model)), name) -end - -function (p::Processor)(graph::FeatureGraph, ps, st::NamedTuple{fields}) where {fields} - process!(p.layers, graph, ps, st) -end - -function process!(layers::NamedTuple{fields}, graph::FeatureGraph, - ps, st::NamedTuple{fields}) where {fields} - uef, ste = update_edge_features( - layers[:edge_model_fn], ps[:edge_model_fn], st[:edge_model_fn], graph) - unf, stn = update_node_features( - layers[:node_model_fn], ps[:node_model_fn], st[:node_model_fn], graph, uef) - new_st = NamedTuple{fields}((stn, ste)) - - return update_features!(graph; nf = graph.nf + unf, ef = graph.ef + uef), new_st -end - -@inline function update_edge_features(el, ps, st, graph::FeatureGraph) - features = aggregate_edge_features(graph) - - return el(features, ps, st) -end - -@inline function update_node_features( - nl, ps, st, graph::FeatureGraph, updated_edge_features) - features = aggregate_node_features(graph, updated_edge_features) - - return nl(features, ps, st) +function (e::Encoder)(graph::FeatureGraph, ps, st) + nf, stn = e.node_layer(graph.nf, ps.node_layer, st.node_layer) + ef, ste = e.edge_layer(graph.ef, ps.edge_layer, st.edge_layer) + return FeatureGraph(graph; nf = nf, ef = ef), (; node_layer = stn, edge_layer = ste) end -struct Decoder{T <: NamedTuple, N <: Lux.NAME_TYPE} <: - Lux.AbstractExplicitContainerLayer{(:layers,)} - layers::T - name::N +struct Processor{N, E} <: Lux.AbstractLuxContainerLayer{(:node_layer, :edge_layer)} + node_layer::N + edge_layer::E end -function Decoder(model; name::Lux.NAME_TYPE = nothing) - fields = (Symbol("model"),) - - return Decoder(NamedTuple{fields}((model,)), name) +function (p::Processor)(graph::FeatureGraph, ps, st) + uef, ste = p.edge_layer(aggregate_edge_features(graph), ps.edge_layer, st.edge_layer) + unf, stn = p.node_layer( + aggregate_node_features(graph, uef), ps.node_layer, st.node_layer) + return FeatureGraph(graph; nf = graph.nf + unf, ef = graph.ef + uef), + (; node_layer = ste, edge_layer = stn) end -function (d::Decoder)(graph::FeatureGraph, ps, st::NamedTuple{fields}) where {fields} - decode!(d.layers, graph, ps, st) +struct Decoder{D} <: Lux.AbstractLuxWrapperLayer{:decode_layer} + decode_layer::D end -function decode!(layers::NamedTuple{fields}, graph::FeatureGraph, - ps, st::NamedTuple{fields}) where {fields} - y, stm = layers[:model](graph.nf, ps[:model], st[:model]) - new_st = NamedTuple{fields}((stm,)) - - return y, new_st +function (d::Decoder)(graph::FeatureGraph, ps, st) + df, std = d.decode_layer(graph.nf, ps, st) + return df, std end diff --git a/src/graph_network.jl b/src/graph_network.jl index a73878f..dcba726 100644 --- a/src/graph_network.jl +++ b/src/graph_network.jl @@ -1,5 +1,6 @@ # # Copyright (c) 2023 Julian Trommer +# Copyright (c) 2026 Josef Jouaux # Licensed under the MIT license. See LICENSE file in the project root for details. # @@ -8,7 +9,8 @@ using ComponentArrays import DataFrames: DataFrame import JLD2: load, save import Statistics: mean -import Zygote: pullback +import Setfield: @set! +import Zygote: withgradient include("feature_graph.jl") include("graph_net_blocks.jl") @@ -27,9 +29,7 @@ The central data structure that contains the neural network and the normalisers - `o_norm`: Normaliser for the output of the GNN, whereas each quantity of interest has its own normaliser. """ mutable struct GraphNetwork - model::Chain - ps::ComponentArray - st::NamedTuple + train_state::Lux.Training.TrainState e_norm::Union{NormaliserOffline, NormaliserOnline} n_norm::Dict{String, Union{NormaliserOffline, NormaliserOnline}} o_norm::Dict{String, Union{NormaliserOffline, NormaliserOnline}} @@ -54,15 +54,14 @@ Constructs a MLP with the given parameters. """ function build_mlp(input_size::T, latent_size::T, output_size::T, hidden_layers::T; layer_norm = true) where {T <: Integer} + mlp = Lux.Chain(Lux.Dense(input_size, latent_size, relu), + collect(Lux.Dense(latent_size, latent_size, relu) for _ in 1:hidden_layers), + Lux.Dense(latent_size, output_size)) if layer_norm - return Chain(Dense(input_size, latent_size, relu), - collect(Dense(latent_size, latent_size, relu) for _ in 1:hidden_layers), - Dense(latent_size, output_size), LayerNorm((output_size,))) - else - return Chain(Dense(input_size, latent_size, relu), - collect(Dense(latent_size, latent_size, relu) for _ in 1:hidden_layers), - Dense(latent_size, output_size)) + mlp = Lux.Chain(mlp.layers..., Lux.LayerNorm((output_size,))) + # mlp = Lux.Chain(mlp.layers..., Lux.LayerNorm((output_size,); dims = 1)) end + return mlp end """ @@ -83,20 +82,22 @@ Constructs the Encode-Process-Decode model as a [Lux.jl](https://github.com/LuxD """ function build_model(quantities_size::Integer, dims, output_size::Integer, mps::Integer, layer_size::Integer, hidden_layers::Integer) - encoder = Encoder(build_mlp(quantities_size, layer_size, layer_size, hidden_layers), + encoder = Encoder( + build_mlp(quantities_size, layer_size, layer_size, hidden_layers), build_mlp(dims + 1, layer_size, layer_size, hidden_layers)) processors = Vector{Processor}() for _ in 1:mps push!(processors, - Processor(build_mlp(2 * layer_size, layer_size, layer_size, hidden_layers), + Processor( + build_mlp(2 * layer_size, layer_size, layer_size, hidden_layers), build_mlp(3 * layer_size, layer_size, layer_size, hidden_layers))) end decoder = Decoder(build_mlp( layer_size, layer_size, output_size, hidden_layers; layer_norm = false)) - model = Chain(encoder, processors..., decoder) + model = Lux.Chain(encoder, processors..., decoder) return model end @@ -117,10 +118,9 @@ Calculates the loss of the network based on the given loss function. ## Returns - Calculated Loss. """ -function loss(ps, gn::GraphNetwork, graph::FeatureGraph, target::AbstractArray{Float32, 2}, +function loss(gn::GraphNetwork, graph::FeatureGraph, target::AbstractArray{Float32, 2}, mask::AbstractArray{T, 1}, loss_function) where {T <: Integer} - output, st = gn.model(graph, ps, gn.st) - gn.st = st + output, _ = gn.train_state.model(graph, gn.train_state.ps, gn.train_state.st) error = loss_function(target, output) @@ -144,14 +144,42 @@ end - Calculated training loss. """ function step!(gn, graph, target_quantities_change, mask, loss_function) - train_loss, back = pullback( - ps -> loss(ps, gn, graph, target_quantities_change, mask, loss_function), gn.ps) - - gs = back(one(train_loss)) + train_loss, + gs = withgradient( + ps -> loss(gn, graph, target_quantities_change, mask, loss_function), + gn.train_state.parameters) return gs, train_loss end +function set_training!(gn::GraphNetwork, training::Bool) + if gn.e_norm isa NormaliserOnline + if training + gn.e_norm.num_accumulations -= gn.e_norm.max_accumulations + else + gn.e_norm.num_accumulations += gn.e_norm.max_accumulations + end + end + for nn in values(gn.n_norm) + if nn isa NormaliserOnline + if training + nn.num_accumulations -= nn.max_accumulations + else + nn.num_accumulations += nn.max_accumulations + end + end + end + for on in values(gn.o_norm) + if on isa NormaliserOnline + if training + on.num_accumulations -= on.max_accumulations + else + on.num_accumulations += on.max_accumulations + end + end + end +end + """ save!(gn, opt_state, df_train, df_valid, step, train_loss, path; is_training = true) @@ -169,20 +197,18 @@ Creates a checkpoint of the [`GraphNetwork`](@ref) at the given training step. ## Keyword Arguments - `is_training = true`: True if used in training, false otherwise (in validation). """ -function save!(gn, opt_state, df_train::DataFrame, df_valid::DataFrame, - step::Integer, train_loss::Float32, path::String; is_training = true) - if is_training - push!(df_train, [step, train_loss]) - else - push!(df_valid, [step, train_loss]) - end +function save!(gn::GraphNetwork, opt_state, df_train::DataFrame, df_valid::DataFrame, + step::Integer, path::String) + ps_data = cpu_device()(getdata(gn.train_state.parameters)) + ps_axes = cpu_device()(getaxes(gn.train_state.parameters)) + st = cpu_device()(gn.train_state.states) save(joinpath(path, "checkpoint_$step.jld2"), - Dict("ps_data" => cpu_device()(getdata(gn.ps)), "ps_axes" => getaxes(gn.ps), - "st" => cpu_device()(gn.st), "e_norm" => serialize(gn.e_norm), + Dict("ps_data" => ps_data, "ps_axes" => ps_axes, + "st" => st, "e_norm" => serialize(gn.e_norm), "n_norm" => serialize(gn.n_norm), "o_norm" => serialize(gn.o_norm), - "opt_state" => cpu_device()(opt_state), - "df_train" => df_train, "df_valid" => df_valid)) + "opt_state" => cpu_device()(opt_state), "df_train" => df_train, + "df_valid" => df_valid)) if isfile(joinpath(path, "checkpoints")) cps = readlines(joinpath(path, "checkpoints")) @@ -232,25 +258,32 @@ function load(quantities, dims, e_norms::Union{NormaliserOffline, NormaliserOnli 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( + 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 + ps = ComponentArray(ps_data, ps_axes) + model = build_model(quantities, dims, output, message_steps, ls, hl) 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) - gn = GraphNetwork(model, ps, st, en, nn, on) + ps = ps |> device + st = st |> device - if !isnothing(opt) - return gn, nothing, df_train, df_valid - else - return gn, device(opt_state), df_train, df_valid - end + train_state = Lux.Training.TrainState(model, ps, st, opt) + @set! train_state.opt_state = opt_state + + gn = GraphNetwork(train_state, en, nn, on) + + return gn, df_train, df_valid else model = build_model(quantities, dims, output, message_steps, ls, hl) ps, st = Lux.setup(Random.default_rng(), model) @@ -258,11 +291,13 @@ function load(quantities, dims, e_norms::Union{NormaliserOffline, NormaliserOnli ps = ComponentArray(ps) |> device st = st |> device - gn = GraphNetwork(model, ps, st, e_norms, n_norms, o_norms) + train_state = Lux.Training.TrainState(model, ps, st, opt) + + gn = GraphNetwork(train_state, e_norms, n_norms, o_norms) df_train = DataFrame(; step = Integer[], loss = Float32[]) df_valid = DataFrame(; step = Integer[], loss = Float32[]) - return gn, nothing, df_train, df_valid + return gn, df_train, df_valid end end diff --git a/src/normaliser.jl b/src/normaliser.jl index 05d7881..5eb296e 100644 --- a/src/normaliser.jl +++ b/src/normaliser.jl @@ -17,19 +17,26 @@ It is recommended to use offline normalization since the minimum and maximum do - `target_min`: Minimum of the target of normalization. - `target_max`: Maximum of the target of normalization. """ -mutable struct NormaliserOfflineMinMax <: NormaliserOffline - data_min::Float32 - data_max::Float32 - target_min::Float32 - target_max::Float32 +mutable struct NormaliserOfflineMinMax{AT, T} <: NormaliserOffline + data_min::AT + data_max::AT + target_min::T + target_max::T +end +function NormaliserOfflineMinMax( + data_min::AT, data_max::AT, device::Function) where {AT} + NormaliserOfflineMinMax(device(data_min), device(data_max), 0.0f0, 1.0f0) end -function NormaliserOfflineMinMax(data_min::Float32, data_max::Float32) - NormaliserOfflineMinMax(data_min, data_max, 0.0f0, 1.0f0) +function NormaliserOfflineMinMax( + data_min::AT, data_max::AT, target_min::T, target_max::T, device::Function) where { + AT, T} + NormaliserOfflineMinMax(device(data_min), device(data_max), target_min, target_max) end -function NormaliserOfflineMinMax(d::Dict{String, Any}) - NormaliserOfflineMinMax(d["data_min"], d["data_max"], d["target_min"], d["target_max"]) +function NormaliserOfflineMinMax(d::Dict{String, Any}, device::Function) + NormaliserOfflineMinMax( + device(d["data_min"]), device(d["data_max"]), d["target_min"], d["target_max"]) end function (n::NormaliserOfflineMinMax)(F) @@ -49,7 +56,9 @@ Inverses the normalised data. - Converted data. """ function inverse_data(n::NormaliserOfflineMinMax, data) - return minmaxnorm(data, n.target_min, n.target_max, n.data_min, n.data_max) + # Since the minmax of the output is not known, we let the decoder handle it + return data + # return minmaxnorm(data, n.target_min, n.target_max, n.data_min, n.data_max) end """ @@ -62,22 +71,26 @@ It is recommended to use offline normalization since the minimum and maximum do - `data_mean`: Mean of the quantity in the dataset. - `data_std`: Standard deviation of the quantity in the dataset. """ -mutable struct NormaliserOfflineMeanStd <: NormaliserOffline - data_mean::Float32 - data_std::Float32 - std_epsilon::Float32 +mutable struct NormaliserOfflineMeanStd{AT, T} <: NormaliserOffline + data_mean::AT + data_std::AT + std_epsilon::T end -function NormaliserOfflineMeanStd(data_mean::Float32, data_std::Float32) - NormaliserOfflineMeanStd(data_mean, data_std, 1.0f-8) +function NormaliserOfflineMeanStd( + data_mean::T, data_std::T, device::Function) where {T} + NormaliserOfflineMeanStd(device(data_mean), device(data_std), eps(eltype(T))) end -function NormaliserOfflineMeanStd(d::Dict{String, Any}) +function NormaliserOfflineMeanStd(d::Dict{String, Any}, device::Function) NormaliserOfflineMeanStd( - d["data_mean"], d["data_std"], haskey(d, "std_epsilon") ? d["std_epsilon"] : 1.0f-8) + device(d["data_mean"]), device(d["data_std"]), + haskey(d, "std_epsilon") ? d["std_epsilon"] : eps(eltype(d["data_mean"]))) end -(n::NormaliserOfflineMeanStd)(F) = (F .- n.data_mean) ./ max(n.data_std, n.std_epsilon) +function (n::NormaliserOfflineMeanStd)(F) + (F .- n.data_mean) ./ max.(n.data_std, n.std_epsilon) +end """ inverse_data(n, data) @@ -92,7 +105,7 @@ Inverses the normalised data. - Converted data. """ function inverse_data(n::NormaliserOfflineMeanStd, data) - return data .* max(n.data_std, n.std_epsilon) .+ n.data_mean + return data .* max.(n.data_std, n.std_epsilon) .+ n.data_mean end """ @@ -109,13 +122,13 @@ It is recommended to use offline normalization since the minimum and maximum do - `acc_sum`: Sum of quantities in each step. - `acc_sum_squared`: Sum of quantities squared in each step. """ -mutable struct NormaliserOnline{T <: AbstractArray{Float32}} - max_accumulations::Float32 - std_epsilon::Float32 - acc_count::Float32 - num_accumulations::Float32 - acc_sum::T - acc_sum_squared::T +mutable struct NormaliserOnline{T, AT} + max_accumulations::T + std_epsilon::T + acc_count::T + num_accumulations::T + acc_sum::AT + acc_sum_squared::AT end """ @@ -132,10 +145,18 @@ It is recommended to use offline normalization since the minimum and maximum do - `max_acc = 10f6`: Maximum number of accumulation steps. - `std_epsilon = 1f-8`: Epsilon for caluclating the standard deviation. """ -function NormaliserOnline( - dim::Integer, device::Function; max_acc::Float32 = 10.0f6, std_ep::Float32 = 1.0f-8) - NormaliserOnline(max_acc, std_ep, 0.0f0, 0.0f0, - device(zeros(Float32, dim)), device(zeros(Float32, dim))) +function NormaliserOnline(::Type{T}, dim::Integer, device::Function; + max_acc::T = 10.0f6, std_ep::T = eps(T)) where {T} + if device == reactant_device() + NormaliserOnline(Reactant.to_rarray(max_acc; track_numbers = true), + Reactant.to_rarray(std_ep; track_numbers = true), + Reactant.to_rarray(zero(T); track_numbers = true), + Reactant.to_rarray(zero(T); track_numbers = true), + device(zeros(T, dim)), device(zeros(T, dim))) + else + NormaliserOnline( + max_acc, std_ep, zero(T), zero(T), device(zeros(T, dim)), device(zeros(T, dim))) + end end """ @@ -149,15 +170,22 @@ It is recommended to use offline normalization since the minimum and maximum do - `device`: Device where the normaliser should be loaded (see [Lux GPU Management](https://lux.csail.mit.edu/dev/manual/gpu_management#gpu-management)). """ function NormaliserOnline(d::Dict{String, Any}, device::Function) - NormaliserOnline(d["max_accumulations"], d["std_epsilon"], d["acc_count"], - d["num_accumulations"], device(d["acc_sum"]), device(d["acc_sum_squared"])) + if device == reactant_device() + NormaliserOnline(Reactant.to_rarray(d["max_accumulations"]; track_numbers = true), + Reactant.to_rarray(d["std_epsilon"]; track_numbers = true), + Reactant.to_rarray(d["acc_count"]; track_numbers = true), + Reactant.to_rarray(d["num_accumulations"]; track_numbers = true), + device(d["acc_sum"]), device(d["acc_sum_squared"])) + else + NormaliserOnline(d["max_accumulations"], d["std_epsilon"], + d["acc_count"], d["num_accumulations"], + device(d["acc_sum"]), device(d["acc_sum_squared"])) + end end -function (n::NormaliserOnline)(F, acc = true::Bool) - if acc - if n.num_accumulations < n.max_accumulations - accumulate_stats!(n, F) - end +function (n::NormaliserOnline)(F::AbstractArray) + @trace if n.num_accumulations < n.max_accumulations + accumulate_stats!(n, F) end return (F .- get_mean(n)) ./ get_std_with_epsilon(n) @@ -186,6 +214,20 @@ function accumulate_stats!(n::NormaliserOnline, F) n.num_accumulations += 1.0f0 end +function accumulate_stats!(n::NormaliserOnline, F::Reactant.TracedRArray) + n.acc_count += size(F)[2] + n.acc_sum += reduce(+, F; dims = 2)[:, 1] + n.acc_sum_squared += reduce(+, F .^ 2; dims = 2)[:, 1] + n.num_accumulations += 1.0f0 +end + +# function accumulate_stats!(n::NormaliserOnline, F::Reactant.TracedRArray) +# @set n.acc_count = n.acc_count + size(F)[2] +# @set n.acc_sum = n.acc_sum + reduce(+, F; dims = 2)[:, 1] +# @set n.acc_sum_squared = n.acc_sum_squared + reduce(+, F .^ 2; dims = 2)[:, 1] +# @set n.num_accumulations = n.num_accumulations + 1.0f0 +# end + function get_mean(n::NormaliserOnline) safe_count = max(n.acc_count, 1.0f0) @@ -200,11 +242,7 @@ function get_std_with_epsilon(n::NormaliserOnline) end function get_sqrt(n) - if n < 0.0f0 - return convert(typeof(n), -Inf32) - else - return sqrt(n) - end + return sqrt(max(0.0f0, n)) end function serialize(ns::Dict{String, Union{NormaliserOffline, NormaliserOnline}}) @@ -218,10 +256,10 @@ end function serialize(n::NormaliserOnline) return Dict{String, Any}( - "max_accumulations" => n.max_accumulations, - "std_epsilon" => n.std_epsilon, - "acc_count" => n.acc_count, - "num_accumulations" => n.num_accumulations, + "max_accumulations" => cpu_device()(n.max_accumulations), + "std_epsilon" => cpu_device()(n.std_epsilon), + "acc_count" => cpu_device()(n.acc_count), + "num_accumulations" => cpu_device()(n.num_accumulations), "acc_sum" => cpu_device()(n.acc_sum), "acc_sum_squared" => cpu_device()(n.acc_sum_squared) ) @@ -229,8 +267,8 @@ end function serialize(n::NormaliserOfflineMinMax) return Dict{String, Any}( - "data_min" => n.data_min, - "data_max" => n.data_max, + "data_min" => cpu_device()(n.data_min), + "data_max" => cpu_device()(n.data_max), "target_min" => n.target_min, "target_max" => n.target_max ) @@ -238,8 +276,8 @@ end function serialize(n::NormaliserOfflineMeanStd) return Dict{String, Any}( - "data_mean" => n.data_mean, - "data_std" => n.data_std + "data_mean" => cpu_device()(n.data_mean), + "data_std" => cpu_device()(n.data_std) ) end @@ -247,9 +285,9 @@ function deserialize(n::Dict{String, Any}, device::Function) if haskey(n, "max_accumulations") return NormaliserOnline(n, device) elseif haskey(n, "data_min") - return NormaliserOfflineMinMax(n) + return NormaliserOfflineMinMax(n, device) elseif haskey(n, "data_mean") - return NormaliserOfflineMeanStd(n) + return NormaliserOfflineMeanStd(n, device) else features = keys(n) norms = deserialize.(values(n), device) diff --git a/src/utils.jl b/src/utils.jl index aa5460c..26bf502 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -47,7 +47,7 @@ end """ one_hot(indices, depth, offset = 0) -Constructs a onehot matrix of Bool with the given indices. +Constructs a onehot matrix of Float32 with the given indices. ## Arguments - `indices`: Indices for the onehot matrix. @@ -58,7 +58,7 @@ Constructs a onehot matrix of Bool with the given indices. - Onehot matrix from the given arguments. """ function one_hot(indices, depth, offset = 0) - result = zeros(Bool, depth, length(indices)) + result = zeros(Float32, depth, length(indices)) for (i, x) in enumerate(indices) if x + offset <= depth && x + offset > 0 result[x + offset, i] = 1 @@ -85,15 +85,16 @@ Normalizes the given input to the new given range. """ function minmaxnorm( input::AbstractArray, input_min, input_max, new_min = 0.0f0, new_max = 1.0f0) - @assert input_min<=input_max "minimum of input has to be lower than or equal to maximum of input : $input_min > $input_max" - @assert new_min<=new_max "minimum of output has to be lower than or equal to maximum of output : $new_min > $new_max" - if input_min == input_max - return typeof(input) <: CuArray ? gpu_device()(zeros(Float32, size(input))) : - zeros(Float32, size(input)) - else - return ((input .- input_min) / (input_max - input_min)) * (new_max - new_min) .+ - new_min - end + @assert minimum(input_max - input_min)>0.0f0 "minimum of input has to be lower than maximum of input : $input_min >= $input_max" + @assert minimum(new_max - new_min)>0.0f0 "minimum of output has to be lower than maximum of output : $new_min >= $new_max" + return ((input .- input_min) ./ (input_max - input_min)) .* (new_max - new_min) .+ + new_min +end + +function minmaxnorm( + input::Reactant.TracedRArray{T, 2}, input_min, input_max, new_min = 0.0f0, new_max = 1.0f0) where {T} + return ((input .- input_min) ./ (input_max - input_min)) .* (new_max - new_min) .+ + new_min end """ @@ -108,7 +109,7 @@ Calculates the mean squared error of the given arguments with [Tullio](https://g ## Returns - Calculated mean squared error. """ -mse_reduce(target, output) = begin +function mse_reduce(target, output) if ndims(target) != 2 || ndims(output) != 2 throw(ArgumentError("Only supported number of dimensions is 2: dims = (target => $(ndims(target)), output => $(ndims(output)))")) end @@ -127,7 +128,7 @@ Implementation of the function [`reducesum`](@ref) with [Tullio](https://github. ## Returns - Reduced array. """ -tullio_reducesum(a, dims) = begin +function tullio_reducesum(a, dims) if dims != 1 && dims != 2 throw(ArgumentError("Only supported dims are 1 and 2: dims = $dims")) end diff --git a/test/runtests.jl b/test/runtests.jl index 55bdc5a..ac10fb0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,6 @@ # # Copyright (c) 2023 Julian Trommer +# Copyright (c) 2026 Josef Jouaux # Licensed under the MIT license. See LICENSE file in the project root for details. # @@ -7,7 +8,8 @@ using GraphNetCore using Test using Aqua -using CUDA, cuDNN, Lux +using CUDA +using Lux import Random: MersenneTwister @@ -61,9 +63,9 @@ import Random: MersenneTwister 0 0 1 0 0 1 0 0] - @test minmaxnorm([2.0f0], 1.0f0, 1.0f0) == [0.0f0] + @test_throws AssertionError minmaxnorm([2.0f0], 1.0f0, 1.0f0) hascuda && - @test minmaxnorm(gpu([1.0f0, 2.0f0]), 1.0f0, 1.0f0) == gpu([0.0f0, 0.0f0]) + @test_throws AssertionError minmaxnorm(gpu([1.0f0, 2.0f0]), 1.0f0, 1.0f0) @test minmaxnorm([1.4f0, 2.3f0, 3.9f0, 4.0f0], -4.0f0, 4.0f0, 0.0f0, 1.0f0) == [0.675f0, 0.7875f0, 0.9875f0, 1.0f0] @test_throws AssertionError minmaxnorm([2.0f0], 1.5f0, 0.5f0) @@ -104,11 +106,13 @@ import Random: MersenneTwister norm_dict_gpu = Dict{String, Union{NormaliserOffline, NormaliserOnline}}( "norm_off" => norm_off, "norm_on" => norm_on_gpu) + # inverse_data for NormaliserOfflineMinMax intentionally returns the input + # unchanged (the decoder handles the output range). @test inverse_data(norm_off, [0.0f0]) == [0.0f0] @test inverse_data(norm_off, [-0.5f0, -0.25f0, 0.1f0, 0.75f0]) == - [-5.0f0, -2.5f0, 1.0f0, 7.5f0] + [-0.5f0, -0.25f0, 0.1f0, 0.75f0] hascuda && @test inverse_data(norm_off, gpu([-0.5f0, -0.25f0, 0.1f0, 0.75f0])) == - gpu([-5.0f0, -2.5f0, 1.0f0, 7.5f0]) + gpu([-0.5f0, -0.25f0, 0.1f0, 0.75f0]) norm_dict_cpu_test = GraphNetCore.deserialize( GraphNetCore.serialize(norm_dict_cpu), cpu) @@ -166,8 +170,8 @@ import Random: MersenneTwister 3.0f0 4.0f0 5.0f0 3.0f0 4.0f0 5.0f0] senders = [2, 3, 3, 1, 1, 2] receivers = [1, 1, 2, 2, 3, 3] - output = [0.7324772f0 -0.027799817f0 0.1475548f0; - 0.42122957f0 -0.6571782f0 -0.15739384f0] + output = [3.0774236f0 -0.37687588f0 0.2191811f0; + -2.30065f0 -2.6680458f0 -1.881568f0] graph = FeatureGraph(nf, ef, senders, receivers)