adjustment to GINConv

Author: CarloLucibello

Project.toml

@@ -13,6 +13,7 @@ KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LearnBase = "7f8f8fb0-2700-5f03-b4bd-41f8cfc144b6"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
@@ -21,6 +22,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+Adapt = "3"
 CUDA = "3.3"
 ChainRulesCore = "1"
 DataStructures = "0.18"

src/GraphNeuralNetworks.jl

@@ -63,5 +63,6 @@ include("msgpass.jl")
 include("layers/basic.jl")
 include("layers/conv.jl")
 include("layers/pool.jl")
+include("deprecations.jl")
 
 end

src/deprecations.jl

@@ -0,0 +1,3 @@
+# Deprecated in v0.1 
+
+@deprecate GINConv(nn; eps=0, aggr=+)  GINConv(nn, eps; aggr)

src/layers/conv.jl

@@ -407,7 +407,7 @@ end
 
 
 @doc raw"""
-    GINConv(f; eps = 0f0)
+    GINConv(f, ϵ; aggr=+)
 
 Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf)
 
@@ -420,30 +420,38 @@ where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer o
 # Arguments
 
 - `f`: A (possibly learnable) function acting on node features. 
-- `eps`: Weighting factor.
+- `ϵ`: Weighting factor.
 """
 struct GINConv{R<:Real} <: GNNLayer
     nn
-    eps::R
+    ϵ::R
+    aggr
 end
 
 @functor GINConv
-Flux.trainable(l::GINConv) = (nn=l.nn,)
+Flux.trainable(l::GINConv) = (l.nn,)
+
+GINConv(nn, ϵ; aggr=+) = GINConv(nn, ϵ, aggr)
 
-function GINConv(nn; eps=0f0)
-    GINConv(nn, eps)
-end
 
 compute_message(l::GINConv, x_i, x_j, e_ij) = x_j 
-update_node(l::GINConv, m, x) = l.nn((1 + l.eps) * x + m)
+update_node(l::GINConv, m, x) = l.nn((1 + ofeltype(x, l.ϵ)) * x + m)
 
 function (l::GINConv)(g::GNNGraph, X::AbstractMatrix)
     check_num_nodes(g, X)
-    X, _ = propagate(l, g, +, X)
+    X, _ = propagate(l, g, l.aggr, X)
     X
 end
 
 
+function Base.show(io::IO, l::GINConv)
+    print(io, "GINConv($(l.nn)")
+    print(io, ", $(l.ϵ)")
+    print(io, ")")
+end
+
+
+
 @doc raw"""
     NNConv(in => out, f, σ=identity; aggr=+, bias=true, init=glorot_uniform)
 

src/utils.jl

@@ -67,3 +67,6 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
     end
     return data
 end
+
+
+ofeltype(x, y) = convert(float(eltype(x)), y)

test/examples/node_classification_cora.jl

@@ -17,11 +17,11 @@ end
 
 # arguments for the `train` function 
 Base.@kwdef mutable struct Args
-    η = 1f-3             # learning rate
-    epochs = 20         # number of epochs
+    η = 5f-3             # learning rate
+    epochs = 10         # number of epochs
     seed = 17             # set seed > 0 for reproducibility
     usecuda = false      # if true use cuda (if available)
-    nhidden = 128        # dimension of hidden features
+    nhidden = 64        # dimension of hidden features
 end
 
 function train(Layer; verbose=false, kws...)
@@ -49,7 +49,7 @@ function train(Layer; verbose=false, kws...)
 
     ## DEFINE MODEL
     model = GNNChain(Layer(nin, nhidden),
-                     Dropout(0.5),
+                    #  Dropout(0.5),
                      Layer(nhidden, nhidden), 
                      Dense(nhidden, nout))  |> device
 
@@ -84,16 +84,16 @@ for Layer in [
             (nin, nout) -> GraphConv(nin => nout, relu, aggr=mean),
             (nin, nout) -> SAGEConv(nin => nout, relu),
             (nin, nout) -> GATConv(nin => nout, relu),
-            (nin, nout) -> GATConv(nin => nout÷2, relu, heads=2),
-            (nin, nout) -> GINConv(Dense(nin, nout, relu)),
-            (nin, nout) -> ChebConv(nin => nout, 3),
-            (nin, nout) -> ResGatedGraphConv(nin => nout, relu),          
+            (nin, nout) -> GINConv(Dense(nin, nout, relu), 0.01, aggr=mean),
+            (nin, nout) -> ChebConv(nin => nout, 2),
+            (nin, nout) -> ResGatedGraphConv(nin => nout, relu),        
             # (nin, nout) -> NNConv(nin => nout),  # needs edge features
             # (nin, nout) -> GatedGraphConv(nout, 2),  # needs nin = nout
             # (nin, nout) -> EdgeConv(Dense(2nin, nout, relu)), # Fits the traning set but does not generalize well
               ]
-    train_res, test_res = train(Layer, verbose=false)
-    # @show Layer(2,2) train_res, test_res
+
+    @show Layer(2,2)
+    train_res, test_res = train(Layer, verbose=true)
     @test train_res.acc > 95
     @test test_res.acc > 70
 end

test/layers/conv.jl

@@ -111,10 +111,10 @@
 
     @testset "GINConv" begin
         nn = Dense(in_channel, out_channel)
-        eps = 0.001f0
-        l = GINConv(nn, eps=eps)
+        
+        l = GINConv(nn, 0.01f0, aggr=mean)
         for g in test_graphs
-            test_layer(l, g, rtol=1e-5, outsize=(out_channel, g.num_nodes), exclude_grad_fields=[:eps]) 
+            test_layer(l, g, rtol=1e-5, outsize=(out_channel, g.num_nodes)) 
         end
 
         @test !in(:eps, Flux.trainable(l))