implement GNNChain

Author: CarloLucibello

Project.toml

@@ -25,6 +25,7 @@ DataStructures = "0.18"
 Flux = "0.12"
 KrylovKit = "0.5"
 LightGraphs = "1.3"
+MacroTools = "0.5"
 NNlib = "0.7"
 NNlibCUDA = "0.1"
 julia = "1.6"

docs/make.jl

@@ -9,10 +9,11 @@ makedocs(;
     pages = ["Home" => "index.md",
              "GNNGraph" => "gnngraph.md",
              "Message Passing" => "messagepassing.md",
-             "Building Models" => "models.md",
+             "Model Building" => "models.md",
              "API Reference" =>
                [
                 "GNNGraph" => "api/gnngraph.md",
+                "Basic Layers" => "api/basic.md",
                 "Convolutional Layers" => "api/conv.md",
                 "Pooling Layers" => "api/pool.md",
                ],

docs/src/api/basic.md

@@ -0,0 +1,12 @@
+# Basic Layers
+
+```@index
+Order = [:type, :function]
+Pages   = ["api/basics.md"]
+```
+
+```@autodocs
+Modules = [GraphNeuralNetworks]
+Pages   = ["layers/basic.jl"]
+Private = false
+```

docs/src/models.md

@@ -1 +1,52 @@
 # Models
+
+## Explicit modeling
+
+```julia
+using Flux, GraphNeuralNetworks
+using Flux: @functor
+
+struct GNN
+    conv1
+    bn
+    conv2
+    dropout
+    dense
+end
+
+@functor GNN
+
+function GNN(din, d, dout)    
+    GNN(GCNConv(din => d),
+        BatchNorm(d),
+        GraphConv(d => d, relu),
+        Dropout(0.5),
+        Dense(d, dout))
+end
+
+function (model::GNN)(g::GNNGraph, x)
+    x = model.conv1(g, x)
+    x = relu.(model.bn(x))
+    x = model.conv2(g, x)
+    x = model.dropout(x)
+    x = model.dense(x)
+    return x 
+end
+
+din, d, dout = 3, 4, 2 
+g = GNNGraph(random_regular_graph(10, 4), graph_type=GRAPH_T)
+X = randn(Float32, din, 10)
+model = GNN(din, d, dout)
+y = model(g, X)
+```
+
+## Compact modeling with GNNChains
+
+```julia
+model = GNNChain(GCNConv(din => d),
+                 BatchNorm(d),
+                 x -> relu.(x),
+                 GraphConv(d => d, relu),
+                 Dropout(0.5),
+                 Dense(d, dout))
+```

examples/cora.jl

@@ -9,28 +9,6 @@ using Statistics, Random
 using CUDA
 CUDA.allowscalar(false)
 
-struct GNN
-    conv1
-    conv2 
-    dense
-end
-
-@functor GNN
-
-function GNN(; nin, nhidden, nout)
-    GNN(GCNConv(nin => nhidden, relu),
-        GCNConv(nhidden => nhidden, relu), 
-        Dense(nhidden, nout))
-end
-
-function (net::GNN)(g, x)
-    x = net.conv1(g, x)
-    x = dropout(x, 0.5)
-    x = net.conv2(g, x)
-    x = net.dense(x)
-    return x
-end
-
 function eval_loss_accuracy(X, y, ids, model, g)
     ŷ = model(g, X)
     l = logitcrossentropy(ŷ[:,ids], y[:,ids])
@@ -63,6 +41,7 @@ function train(; kws...)
         @info "Training on CPU"
     end
 
+    # LOAD DATA
     data = Cora.dataset()
     g = GNNGraph(data.adjacency_list) |> device
     X = data.node_features |> device
@@ -72,14 +51,20 @@ function train(; kws...)
     test_ids = data.test_indices |> device
     ytrain = y[:,train_ids]
 
-    model = GNN(nin=size(X,1), 
-                nhidden=args.nhidden, 
-                nout=data.num_classes) |> device
+    nin, nhidden, nout = size(X,1), args.nhidden, data.num_classes 
+    
+    ## DEFINE MODEL
+    model = GNNGraph(GCNConv(nin => nhidden, relu),
+                     Dropout(0.5)
+                     GCNConv(nhidden => nhidden, relu), 
+                     Dense(nhidden, nout))  |> device
+
     ps = Flux.params(model)
     opt = ADAM(args.η)
 
     @info "NUM NODES: $(g.num_nodes)  NUM EDGES: $(g.num_edges)"
 
+    ## LOGGING FUNCTION
     function report(epoch)
         train = eval_loss_accuracy(X, y, train_ids, model, g)
         test = eval_loss_accuracy(X, y, test_ids, model, g)        

src/GraphNeuralNetworks.jl

@@ -1,11 +1,13 @@
 module GraphNeuralNetworks
 
+using Core: apply_type
 using NNlib: similar
 using LinearAlgebra: similar, fill!
 using Statistics: mean
 using LinearAlgebra
 using SparseArrays
 import KrylovKit
+using Base: tail
 using CUDA
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
@@ -27,7 +29,8 @@ export
     # from LightGraphs
     adjacency_matrix, 
 
-    # layers/msgpass
+    # msgpass
+    # update, update_edge, update_global, message, propagate,
 
     # layers/basic
     GNNLayer,
@@ -54,11 +57,9 @@ export
 include("gnngraph.jl")
 include("graph_conversions.jl")
 include("utils.jl")
-
-include("layers/msgpass.jl")
+include("msgpass.jl")
 include("layers/basic.jl")
 include("layers/conv.jl")
 include("layers/pool.jl")
 
-
 end

src/layers/basic.jl

@@ -52,12 +52,15 @@ end
 @forward GNNChain.layers Base.getindex, Base.length, Base.first, Base.last,
     Base.iterate, Base.lastindex, Base.keys
 
-functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)
+Flux.functor(::Type{<:GNNChain}, c) = c.layers, ls -> GNNChain(ls...)
 
-applychain(::Tuple{}, x) = x
-applychain(fs::Tuple, x) = applychain(tail(fs), first(fs)(x))
+applylayer(l, g::GNNGraph, x) = l(x)
+applylayer(l::GNNLayer, g::GNNGraph, x) = l(g, x)
 
-(c::GNNChain)(x) = applychain(Tuple(c.layers), x)
+applychain(::Tuple{}, g::GNNGraph, x) = x
+applychain(fs::Tuple, g::GNNGraph, x) = applychain(tail(fs), g, applylayer(first(fs), g, x))
+
+(c::GNNChain)(g::GNNGraph, x) = applychain(Tuple(c.layers), g, x)
 
 Base.getindex(c::GNNChain, i::AbstractArray) = GNNChain(c.layers[i]...)
 Base.getindex(c::GNNChain{<:NamedTuple}, i::AbstractArray) = 

src/layers/msgpass.jl renamed to src/msgpass.jl

@@ -1,3 +1,34 @@
 @testset "basic" begin
+    @testset "GNNChain" begin
+        n, din, d, dout = 10, 3, 4, 2
+        
+        g = GNNGraph(random_regular_graph(n, 4), graph_type=GRAPH_T)
+        
+        gnn = GNNChain(GCNConv(din => d),
+                       BatchNorm(d),
+                       x -> relu.(x),
+                       GraphConv(d => d, relu),
+                       Dropout(0.5),
+                       Dense(d, dout))
+        
+        X = randn(Float32, din, n)
 
+        y = gnn(g, X)
+  
+        @test y isa Matrix{Float32}
+        @test size(y) == (dout, n)
+
+        @test length(params(gnn)) == 9
+        
+        gs = gradient(x -> sum(gnn(g, x)), X)[1]
+        @test gs isa Matrix{Float32}
+        @test size(gs) == size(X) 
+
+        gs = gradient(() -> sum(gnn(g, X)), Flux.params(gnn))
+        for p in Flux.params(gnn)
+            @test eltype(gs[p]) == Float32
+            @test size(gs[p]) == size(p)
+        end
+    end
 end
+