Merge pull request #15 from CarloLucibello/cl/gnnchain

add GNNLayer, GNNChain, GNNGraph

Author: CarloLucibello

Project.toml

@@ -11,6 +11,8 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 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"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -24,6 +26,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/Project.toml

@@ -1,3 +1,5 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"

docs/make.jl

@@ -1,4 +1,4 @@
-using GraphNeuralNetworks
+using Flux, NNlib, GraphNeuralNetworks
 using Documenter
 
 DocMeta.setdocmeta!(GraphNeuralNetworks, :DocTestSetup, :(using GraphNeuralNetworks); recursive=true)
@@ -7,14 +7,17 @@ makedocs(;
     modules=[GraphNeuralNetworks],
     sitename = "GraphNeuralNetworks.jl",
     pages = ["Home" => "index.md",
-             "Graphs" => "graphs.md",
-             "Message passing" => "messagepassing.md",
-             "Building models" => "models.md",
+             "GNNGraph" => "gnngraph.md",
+             "Message Passing" => "messagepassing.md",
+             "Model Building" => "models.md",
              "API Reference" =>
                [
-                "Graphs" => "api/graphs.md",
+                "GNNGraph" => "api/gnngraph.md",
+                "Basic Layers" => "api/basic.md",
                 "Convolutional Layers" => "api/conv.md",
                 "Pooling Layers" => "api/pool.md",
+                "Message Passing" => "api/messagepassing.md",
+                "NNlib" => "api/nnlib.md",
                ],
               "Developer Notes" => "dev.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/api/graphs.md renamed to docs/src/api/gnngraph.md

@@ -7,6 +7,6 @@ Pages   = ["api/graphs.md"]
 
 ```@autodocs
 Modules = [GraphNeuralNetworks]
-Pages   = ["featuredgraph.jl"]
+Pages   = ["gnngraph.jl"]
 Private = false
 ```

docs/src/api/messagepassing.md

@@ -0,0 +1,7 @@
+# Message Passing
+
+```@docs
+GraphNeuralNetworks.message
+GraphNeuralNetworks.update
+GraphNeuralNetworks.propagate
+```

docs/src/api/nnlib.md

@@ -0,0 +1,10 @@
+# NNlib
+
+Primitive functions implemented in NNlib.jl.
+
+```@docs
+NNlib.gather!
+NNlib.gather
+NNlib.scatter!
+NNlib.scatter
+```

docs/src/graphs.md renamed to docs/src/gnngraph.md

@@ -1 +1,88 @@
 # Models
+
+GraphNeuralNetworks.jl provides common graph convolutional layers by which you can assemble arbitrarily deep or complex models. GNN layers are compatible with 
+Flux.jl ones, therefore expert Flux's users should be immediately able to define and train 
+their models. 
+
+In what follows, we discuss two different styles for model creation:
+the *explicit modeling* style, more verbose but more flexible, 
+and the *implicit modeling* style based on [`GNNChain`](@ref), more concise but less flexible.
+
+## Explicit modeling
+
+In the explicit modeling style, the model is created according to the following steps:
+
+1. Define a new type for your model (`GNN` in the example below). Layers and submodels are fields.
+2. Apply `Flux.@functor` to the new type to make it Flux's compatible (parameters' collection, gpu movement, etc...)
+3. Optionally define a convenience constructor for your model.
+4. Define the forward pass by implementing the function call method for your type
+5. Instantiate the model. 
+
+Here is an example of this construction:
+```julia
+using Flux, LightGraphs, GraphNeuralNetworks
+using Flux: @functor
+
+struct GNN                                # step 1
+    conv1
+    bn
+    conv2
+    dropout
+    dense
+end
+
+@functor GNN                              # step 2
+
+function GNN(din::Int, d::Int, dout::Int) # step 3    
+    GNN(GCNConv(din => d),
+        BatchNorm(d),
+        GraphConv(d => d, relu),
+        Dropout(0.5),
+        Dense(d, dout))
+end
+
+function (model::GNN)(g::GNNGraph, x)     # step 4
+    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))
+X = randn(Float32, din, 10)
+model = GNN(din, d, dout)                 # step 5
+y = model(g, X)
+```
+
+## Implicit modeling with GNNChains
+
+While very flexible, the way in which we defined `GNN` model definition in last section is a bit verbose.
+In order to simplify things, we provide the [`GNNChain`](@ref) type. It is very similar 
+to Flux's well known `Chain`. It allows to compose layers in a sequential fashion as Chain
+does, propagating the output of each layer to the next one. In addition, `GNNChain` 
+handles propagates the input graph as well, providing it as a first argument
+to layers subtyping the [`GNNLayer`](@ref) abstract type. 
+
+Using `GNNChain`, the previous example becomes
+
+```julia
+using Flux, LightGraphs, GraphNeuralNetworks
+
+din, d, dout = 3, 4, 2 
+g = GNNGraph(random_regular_graph(10, 4))
+X = randn(Float32, din, 10)
+
+model = GNNChain(GCNConv(din => d),
+                 BatchNorm(d),
+                 x -> relu.(x),
+                 GraphConv(d => d, relu),
+                 Dropout(0.5),
+                 Dense(d, dout))
+
+y = model(g, X)
+```
+
+The `GNNChain` only propagates the graph and the node features. More complex scenarios, e.g. when also edge features are updated, have to be handled using the explicit definition of the forward pass. 

docs/src/models.md

@@ -9,30 +9,8 @@ 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)(fg, x)
-    x = net.conv1(fg, x)
-    x = dropout(x, 0.5)
-    x = net.conv2(fg, x)
-    x = net.dense(x)
-    return x
-end
-
-function eval_loss_accuracy(X, y, ids, model, fg)
-    ŷ = model(fg, X)
+function eval_loss_accuracy(X, y, ids, model, g)
+    ŷ = model(g, X)
     l = logitcrossentropy(ŷ[:,ids], y[:,ids])
     acc = mean(onecold(ŷ[:,ids] |> cpu) .== onecold(y[:,ids] |> cpu))
     return (loss = round(l, digits=4), acc = round(acc*100, digits=2))
@@ -63,34 +41,41 @@ function train(; kws...)
         @info "Training on CPU"
     end
 
+    # LOAD DATA
     data = Cora.dataset()
-    fg = FeaturedGraph(data.adjacency_list) |> device
+    g = GNNGraph(data.adjacency_list) |> device
     X = data.node_features |> device
     y = onehotbatch(data.node_labels, 1:data.num_classes) |> device
     train_ids = data.train_indices |> device
     val_ids = data.val_indices |> device
     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 = GNNChain(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: $(fg.num_nodes)  NUM EDGES: $(fg.num_edges)"
+    @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, fg)
-        test = eval_loss_accuracy(X, y, test_ids, model, fg)        
+        train = eval_loss_accuracy(X, y, train_ids, model, g)
+        test = eval_loss_accuracy(X, y, test_ids, model, g)        
         println("Epoch: $epoch   Train: $(train)   Test: $(test)")
     end
 
     ## TRAINING
     report(0)
     for epoch in 1:args.epochs
         gs = Flux.gradient(ps) do
-            ŷ = model(fg, X)
+            ŷ = model(g, X)
             logitcrossentropy(ŷ[:,train_ids], ytrain)
         end