diff --git a/GNNLux/src/layers/conv.jl b/GNNLux/src/layers/conv.jl index f67e54381..73a78f223 100644 --- a/GNNLux/src/layers/conv.jl +++ b/GNNLux/src/layers/conv.jl @@ -1287,7 +1287,7 @@ function Base.show(io::IO, l::GatedGraphConv) end @doc raw""" - GINConv(f, ϵ; aggr=+) + GINConv(f, [ϵ]; aggr=+, train_eps=false) Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf). @@ -1300,7 +1300,9 @@ where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer o # Arguments - `f`: A (possibly learnable) function acting on node features. -- `ϵ`: Weighting factor. +- `ϵ`: Initial value of the weighting factor ``\epsilon``. Default `0.0f0`. +- `train_eps`: If `true`, ``\epsilon`` is trainable. Default `false`. +- `aggr`: Neighborhood aggregation function. Default `+`. # Examples: @@ -1322,7 +1324,7 @@ x = randn(rng, Float32, in_channel, g.num_nodes) nn = Dense(in_channel, out_channel) # create layer -l = GINConv(nn, 0.01f0, aggr = mean) +l = GINConv(nn; ϵ = 0.01f0, train_eps = true, aggr = mean) # setup layer ps, st = LuxCore.setup(rng, l) @@ -1333,15 +1335,31 @@ y, st = l(g, x, ps, st) # size: out_channel × num_nodes """ @concrete struct GINConv <: GNNContainerLayer{(:nn,)} nn <: AbstractLuxLayer - ϵ <: Real + ϵ aggr + train_eps::Bool end -GINConv(nn, ϵ; aggr = +) = GINConv(nn, ϵ, aggr) +GINConv(nn, ϵ::Real; aggr = +, train_eps::Bool = false) = + GINConv(nn, train_eps ? [ϵ] : ϵ, aggr, train_eps) + +GINConv(nn, ϵ::Real, aggr) = GINConv(nn, ϵ; aggr) + +GINConv(nn; ϵ::Real = 0.0f0, aggr = +, train_eps::Bool = false) = + GINConv(nn, ϵ; aggr, train_eps) + +function LuxCore.initialparameters(rng::AbstractRNG, l::GINConv) + nn = LuxCore.initialparameters(rng, l.nn) + return l.train_eps ? (; nn, ϵ = l.ϵ) : (; nn) +end + +LuxCore.parameterlength(l::GINConv) = parameterlength(l.nn) + l.train_eps +LuxCore.statelength(l::GINConv) = statelength(l.nn) function (l::GINConv)(g, x, ps, st) nn = StatefulLuxLayer{true}(l.nn, ps.nn, st.nn) - m = (; nn, l.ϵ, l.aggr) + ϵ = l.train_eps ? ps.ϵ : l.ϵ + m = (; nn, ϵ, l.aggr) y = GNNlib.gin_conv(m, g, x) stnew = (; nn = _getstate(nn)) return y, stnew @@ -1349,7 +1367,9 @@ end function Base.show(io::IO, l::GINConv) print(io, "GINConv($(l.nn)") - print(io, ", $(l.ϵ)") + print(io, ", ϵ=$(l.ϵ)") + l.train_eps && print(io, ", train_eps=true") + l.aggr == (+) || print(io, ", aggr=$(l.aggr)") print(io, ")") end @@ -1879,4 +1899,4 @@ function (l::SAGEConv)(g, x, ps, st) m = (; ps.weight, bias = _getbias(ps), l.σ, l.aggr) return GNNlib.sage_conv(m, g, x), st -end \ No newline at end of file +end diff --git a/GNNLux/test/layers/conv.jl b/GNNLux/test/layers/conv.jl index 5bfb01825..d32c2772e 100644 --- a/GNNLux/test/layers/conv.jl +++ b/GNNLux/test/layers/conv.jl @@ -92,8 +92,17 @@ @testset "GINConv" begin nn = Chain(Dense(in_dims => out_dims, tanh), Dense(out_dims => out_dims)) - l = GINConv(nn, 0.5) + l = GINConv(nn; ϵ = 0.5) test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true) + ps = LuxCore.initialparameters(rng, l) + @test !hasproperty(ps, :ϵ) + + l = GINConv(nn; train_eps = true) + test_lux_layer(rng, l, g, x, sizey=(out_dims,g.num_nodes), container=true) + ps = LuxCore.initialparameters(rng, l) + @test ps.ϵ == [0.0f0] + + @test GINConv(nn, 0.5).ϵ == 0.5 end @testset "MEGNetConv" begin diff --git a/GNNlib/src/layers/conv.jl b/GNNlib/src/layers/conv.jl index 8ae67f5a0..c3b1b77cc 100644 --- a/GNNlib/src/layers/conv.jl +++ b/GNNlib/src/layers/conv.jl @@ -252,7 +252,8 @@ function gin_conv(l, g::AbstractGNNGraph, x) xj, xi = expand_srcdst(g, x) m = propagate(copy_xj, g, l.aggr, xj = xj) - return l.nn((1 .+ ofeltype(xi, l.ϵ)) .* xi .+ m) + ϵ = l.ϵ isa Real ? ofeltype(xi, l.ϵ) : l.ϵ + return l.nn((1 .+ ϵ) .* xi .+ m) end ####################### NNConv ###################################### diff --git a/GraphNeuralNetworks/src/layers/conv.jl b/GraphNeuralNetworks/src/layers/conv.jl index ab038ea01..0af8668e9 100644 --- a/GraphNeuralNetworks/src/layers/conv.jl +++ b/GraphNeuralNetworks/src/layers/conv.jl @@ -590,7 +590,7 @@ function Base.show(io::IO, l::EdgeConv) end @doc raw""" - GINConv(f, ϵ; aggr=+) + GINConv(f, [ϵ]; aggr=+, train_eps=false) Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf). @@ -603,7 +603,9 @@ where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer o # Arguments - `f`: A (possibly learnable) function acting on node features. -- `ϵ`: Weighting factor. +- `ϵ`: Initial value of the weighting factor ``\epsilon``. Default `0.0f0`. +- `train_eps`: If `true`, ``\epsilon`` is trainable. Default `false`. +- `aggr`: Neighborhood aggregation function. Default `+`. # Examples: @@ -619,28 +621,37 @@ g = GNNGraph(s, t) nn = Dense(in_channel, out_channel) # create layer -l = GINConv(nn, 0.01f0, aggr = mean) +l = GINConv(nn; ϵ = 0.01f0, train_eps = true, aggr = mean) # forward pass y = l(g, x) ``` """ -struct GINConv{R <: Real, NN, A} <: GNNLayer +struct GINConv{E, NN, A} <: GNNLayer nn::NN - ϵ::R + ϵ::E aggr::A + train_eps::Bool end Flux.@layer :expand GINConv -Flux.trainable(l::GINConv) = (nn = l.nn,) +Flux.trainable(l::GINConv) = l.train_eps ? (; l.nn, l.ϵ) : (; l.nn) -GINConv(nn, ϵ; aggr = +) = GINConv(nn, ϵ, aggr) +GINConv(nn, ϵ::Real; aggr = +, train_eps::Bool = false) = + GINConv(nn, train_eps ? [ϵ] : ϵ, aggr, train_eps) + +GINConv(nn, ϵ::Real, aggr) = GINConv(nn, ϵ; aggr) + +GINConv(nn; ϵ::Real = 0.0f0, aggr = +, train_eps::Bool = false) = + GINConv(nn, ϵ; aggr, train_eps) (l::GINConv)(g, x) = GNNlib.gin_conv(l, g, x) function Base.show(io::IO, l::GINConv) print(io, "GINConv($(l.nn)") - print(io, ", $(l.ϵ)") + print(io, ", ϵ=$(l.ϵ)") + l.train_eps && print(io, ", train_eps=true") + l.aggr == (+) || print(io, ", aggr=$(l.aggr)") print(io, ")") end diff --git a/GraphNeuralNetworks/test/layers/conv.jl b/GraphNeuralNetworks/test/layers/conv.jl index 7585b817d..7cecaf143 100644 --- a/GraphNeuralNetworks/test/layers/conv.jl +++ b/GraphNeuralNetworks/test/layers/conv.jl @@ -278,19 +278,31 @@ end using .TestModule nn = Dense(D_IN, D_OUT) - l = GINConv(nn, 0.01, aggr = mean) + l = GINConv(nn; ϵ = 0.01, aggr = mean) for g in TEST_GRAPHS @test size(l(g, g.x)) == (D_OUT, g.num_nodes) test_gradients(l, g, g.x, rtol = RTOL_HIGH) end + @test keys(Flux.trainable(l)) == (:nn,) - @test !in(:eps, Flux.trainable(l)) + l = GINConv(Dense(D_IN, D_OUT); train_eps = true, aggr = mean) + @test keys(Flux.trainable(l)) == (:nn, :ϵ) + @test l.ϵ == [0.0f0] + for g in TEST_GRAPHS + @test size(l(g, g.x)) == (D_OUT, g.num_nodes) + end + g = first(TEST_GRAPHS) + grad = Flux.gradient(l -> mean(l(g, g.x)), l)[1] + @test grad.ϵ isa AbstractVector + @test length(grad.ϵ) == 1 + + @test GINConv(nn, 0.5).ϵ == 0.5 end @testitem "GINConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin using .TestModule nn = Dense(D_IN, D_OUT) - l = GINConv(nn, 0.01, aggr = mean) + l = GINConv(nn; train_eps = true, aggr = mean) for g in TEST_GRAPHS g.graph isa AbstractSparseMatrix && continue @test size(l(g, g.x)) == (D_OUT, g.num_nodes)