Merge pull request #30 from CarloLucibello/cl/dev

mode docs

Author: CarloLucibello

README.md

@@ -21,4 +21,4 @@ Its most relevant features are:
 
 ## Usage
 
-Usage examples can be found in the `examples/` folder. 
+Usage examples can be found in the [examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) folder. 

docs/make.jl

@@ -8,7 +8,7 @@ makedocs(;
     doctest=false, clean=true,     
     sitename = "GraphNeuralNetworks.jl",
     pages = ["Home" => "index.md",
-             "GNNGraph" => "gnngraph.md",
+             "Graphs" => "gnngraph.md",
              "Message Passing" => "messagepassing.md",
              "Model Building" => "models.md",
              "API Reference" =>

docs/src/index.md

@@ -11,26 +11,84 @@ Its most relevant features are:
 * Makes it easy to define custom graph convolutional layers.
 
 
+## Package overview
 
+Let's give a brief overview of the package solving a  
+graph regression problem on fake data. 
 
-## Package overview
+Usage examples on real datasets can be found in the [examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) folder. 
 
 ### Data preparation
 
+First, we create our dataset consisting in multiple random graphs and associated data features. 
+that we batch together into a unique graph.
 
-```
-using LightGraphs
+```juliarepl
+julia> using GraphNeuralNetworks, LightGraphs, Flux, CUDA, Statistics
 
-lg = LightGraphs.Graph(5) # create a light's graph graph
-add_edge!(g, 1, 2)
-add_edge!(g, 1, 3)
-add_edge!(g, 2, 4)
-add_edge!(g, 2, 5)
-add_edge!(g, 3, 4)
+julia> all_graphs = GNNGraph[];
 
-g = GNNGraph(g)
+julia> for _ in 1:1000
+           g = GNNGraph(random_regular_graph(10, 4),  
+                       ndata=(; x = randn(Float32, 16,10)),  # input node features
+                       gdata=(; y = randn(Float32)))         # regression target   
+           push!(all_graphs, g)
+       end
+
+julia> gbatch = Flux.batch(all_graphs)
+GNNGraph:
+    num_nodes = 10000
+    num_edges = 20000
+    num_graphs = 1000
+    ndata:
+        x => (16, 10000)
+    edata:
+    gdata:
+        y => (1000,)
 ```
+
+
 ### Model building 
 
+We concisely define our model using as a [`GNNChain`](@ref) containing 2 graph convolutaional 
+layers. If CUDA is available, our model will leave on the gpu.
+
+```juliarepl
+julia> device = CUDA.functional() ? Flux.gpu : Flux.cpu;
+
+julia> model = GNNChain(GCNConv(16 => 64),
+                        BatchNorm(64),
+                        x -> relu.(x),
+                        GCNConv(64 => 64, relu),
+                        GlobalPool(mean),
+                        Dense(64, 1)) |> device;
+
+julia> ps = Flux.params(model);
+
+julia> opt = ADAM(1f-4);
+```
+
 ### Training 
 
+```juliarepl
+gtrain, _ = getgraph(gbatch, 1:800)
+gtest, _ = getgraph(gbatch, 801:gbatch.num_graphs)
+train_loader = Flux.Data.DataLoader(gtrain, batchsize=32, shuffle=true)
+test_loader = Flux.Data.DataLoader(gtest, batchsize=32, shuffle=false)
+
+function loss(g::GNNGraph)
+    mean((vec(model(g, g.ndata.x)) - g.gdata.y).^2)
+end
+
+loss(loader) = mean(loss(g |> device) for g in loader)
+
+for epoch in 1:100
+    for g in train_loader
+        g = g |> gpu
+        grad = gradient(() -> loss(g), ps)
+        Flux.Optimise.update!(opt, ps, grad)
+    end
+
+    @info (; epoch, train_loss=loss(train_loader), test_loss=loss(test_loader))
+end
+```

docs/src/messagepassing.md

@@ -18,4 +18,47 @@ where ``\phi`` is expressed by the [`compute_message`](@ref) function,
 ``\gamma_x`` and ``\gamma_e`` by [`update_node`](@ref) and [`update_edge`](@ref)
 respectively.
 
-See [`GraphConv`](ref) and [`GATConv`](ref)'s implementations as usage examples. 
+## An example: implementing the GCNConv
+
+Let's (re-)implement the [`GCNConv`](@ref) layer use the message passing framework.
+The convolution reads 
+```math
+
+```math
+\mathbf{x}'_i = \sum_{j \in {i} \cup N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
+```
+where ``c_{ij} = \sqrt{(1+|N(i)|)(1+|N(j)|)}``. We will also add a bias and an activation function.
+
+```julia
+using Flux, LightGraphs, GraphNeuralNetworks
+import GraphNeuralNetworks: compute_message, update_node, propagate
+
+struct GCN{A<:AbstractMatrix, B, F} <: GNNLayer
+    weight::A
+    bias::B
+    σ::F
+end
+
+Flux.@functor GCN # allow collecting params, gpu movement, etc...
+
+function GCN(ch::Pair{Int,Int}, σ=identity)
+    in, out = ch
+    W = Flux.glorot_uniform(out, in)
+    b = zeros(Float32, out)
+    GCN(W, b, σ)
+end
+
+compute_message(l::GCN, xi, xj, eij) = l.weight * xj
+update_node(l::GCN, m, x) = m
+
+function (l::GCN)(g::GNNGraph, x::AbstractMatrix{T}) where T
+    g = add_self_loops(g)
+    c = 1 ./ sqrt.(degree(g, T, dir=:in))
+    x = x .* c'
+    x, _ = propagate(l, g, +, x)
+    x = x .* c'
+    return l.σ.(x .+ l.bias)
+end
+```
+
+

src/gnngraph.jl

@@ -80,13 +80,13 @@ g = GNNGraph(s, t)
 g = GNNGraph(erdos_renyi(100, 20))
 
 # Add 2 node feature arrays
-g = GNNGraph(g, ndata = (X = rand(100, g.num_nodes), y = rand(g.num_nodes)))
+g = GNNGraph(g, ndata = (x=rand(100, g.num_nodes), y=rand(g.num_nodes)))
 
-# Add node features and edge features with default names `X` and `E` 
+# Add node features and edge features with default names `x` and `e` 
 g = GNNGraph(g, ndata = rand(100, g.num_nodes), edata = rand(16, g.num_edges))
 
-g.ndata.X
-g.ndata.E
+g.ndata.x
+g.ndata.e
 
 # Send to gpu
 g = g |> gpu
@@ -132,9 +132,9 @@ function GNNGraph(data;
 
     num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1
 
-    ndata = normalize_graphdata(ndata, :X)
-    edata = normalize_graphdata(edata, :E)
-    gdata = normalize_graphdata(gdata, :U)
+    ndata = normalize_graphdata(ndata, :x)
+    edata = normalize_graphdata(edata, :e)
+    gdata = normalize_graphdata(gdata, :u)
 
     GNNGraph(g, 
             num_nodes, num_edges, num_graphs, 
@@ -170,8 +170,7 @@ function Base.show(io::IO, g::GNNGraph)
     println(io, "GNNGraph:
     num_nodes = $(g.num_nodes)
     num_edges = $(g.num_edges)
-    num_graphs = $(g.num_graphs)
-    # feature name => array size")
+    num_graphs = $(g.num_graphs)")
     println(io, "    ndata:")
     for k in keys(g.ndata)
         println(io, "        $k => $(size(g.ndata[k]))")
@@ -498,7 +497,7 @@ function node_features(g::GNNGraph)
     if isempty(g.ndata)
         return nothing
     elseif length(g.ndata) > 1
-        @error "Multiple feature arrays, access directly with g.ndata.X"
+        @error "Multiple feature arrays, access directly through `g.ndata`"
     else
         return g.ndata[1]
     end
@@ -508,7 +507,7 @@ function edge_features(g::GNNGraph)
     if isempty(g.edata)
         return nothing
     elseif length(g.edata) > 1
-        @error "Multiple feature arrays, access directly with g.edata.E"
+        @error "Multiple feature arrays, access directly through `g.edata`"
     else
         return g.edata[1]
     end
@@ -518,7 +517,7 @@ function graph_features(g::GNNGraph)
     if isempty(g.gdata)
         return nothing
     elseif length(g.gdata) > 1
-        @error "Multiple feature arrays, access directly with g.gdata.U"
+        @error "Multiple feature arrays, access directly through `g.gdata`"
     else
         return g.gdata[1]
     end

src/layers/conv.jl

@@ -5,9 +5,9 @@ Graph convolutional layer from paper [Semi-supervised Classification with Graph
 
 Performs the operation
 ```math
-\mathbf{x}'_i = \sum_{j\in N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
+\mathbf{x}'_i = \sum_{j\in \{i\} \cup N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
 ```
-where ``c_{ij} = \sqrt{N(i)\,N(j)}``.
+where ``c_{ij} = \sqrt{(1+|N(i)|)(1+|N(j)|)}``.
 
 The input to the layer is a node feature array `X` 
 of size `(num_features, num_nodes)`.

src/layers/pool.jl

@@ -12,7 +12,7 @@ and performs the operation
 ```
 where ``V`` is the set of nodes of the input graph and 
 the type of aggregation represented by ``\square`` is selected by the `aggr` argument. 
-Commonly used aggregations are are `mean`, `max`, and `+`.
+Commonly used aggregations are `mean`, `max`, and `+`.
 
 ```julia
 using Flux, GraphNeuralNetworks, LightGraphs

src/utils.jl

@@ -12,6 +12,9 @@ end
 
 cat_features(x1::Nothing, x2::Nothing) = nothing 
 cat_features(x1::AbstractArray, x2::AbstractArray) = cat(x1, x2, dims=ndims(x1))
+cat_features(x1::Union{Number, AbstractVector}, x2::Union{Number, AbstractVector}) = 
+    cat(x1, x2, dims=1)
+
 
 function cat_features(x1::NamedTuple, x2::NamedTuple)
     sort(collect(keys(x1))) == sort(collect(keys(x2))) ||

test/gnngraph.jl

@@ -118,6 +118,13 @@
         @test s == [edge_index(g1)[1]; 10 .+ edge_index(g2)[1]; 14 .+ edge_index(g3)[1]] 
         @test t == [edge_index(g1)[2]; 10 .+ edge_index(g2)[2]; 14 .+ edge_index(g3)[2]] 
         @test node_features(g123)[:,11:14] ≈ node_features(g2) 
+
+        # scalar graph features
+        g1 = GNNGraph(random_regular_graph(10,2), gdata=rand())
+        g2 = GNNGraph(random_regular_graph(4,2), gdata=rand())
+        g3 = GNNGraph(random_regular_graph(4,2), gdata=rand())
+        g123 = Flux.batch([g1, g2, g3])
+        @test g123.gdata.u == [g1.gdata.u, g2.gdata.u, g3.gdata.u]
     end
 
     @testset "getgraph"  begin

test/layers/conv.jl

@@ -38,7 +38,7 @@
         @test size(node_features(gt_)) == (out_channel, N)
 
         gs = Zygote.gradient(x -> sum(node_features(l(x))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), l)[1]
         @test size(gs.weight) == size(l.weight)
@@ -74,7 +74,7 @@
         @test size(node_features(gt_)) == (out_channel, N)
 
         gs = Zygote.gradient(x -> sum(node_features(l(x))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), l)[1]
         @test size(gs.weight) == size(l.weight)
@@ -109,7 +109,7 @@
         @test size(node_features(gt_)) == (out_channel, N)
 
         gs = Zygote.gradient(g -> sum(node_features(l(g))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), l)[1]
         @test size(gs.weight1) == size(l.weight1)
@@ -148,7 +148,7 @@
             @test size(node_features(gt_)) == (concat ? (out_channel*heads, N) : (out_channel, N))
 
             gs = Zygote.gradient(g -> sum(node_features(gat(g))), g_gat)[1]
-            @test size(gs.ndata.X) == size(X)
+            @test size(gs.ndata.x) == size(X)
 
             gs = Zygote.gradient(model -> sum(node_features(model(g_gat))), gat)[1]
             @test size(gs.weight) == size(gat.weight)
@@ -182,7 +182,7 @@
         @test size(node_features(gt_)) == (out_channel, N)
 
         gs = Zygote.gradient(x -> sum(node_features(ggc(x))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), ggc)[1]
         @test size(gs.weight) == size(ggc.weight)
@@ -206,7 +206,7 @@
         @test size(node_features(gt_)) == (out_channel, N)
 
         gs = Zygote.gradient(x -> sum(node_features(ec(x))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), ec)[1]
         @test size(gs.nn.weight) == size(ec.nn.weight)
@@ -226,7 +226,7 @@
         @test size(node_features(g_)) == (out_channel, N)
 
         gs = Zygote.gradient(g -> sum(node_features(l(g))), g)[1]
-        @test size(gs.ndata.X) == size(X)
+        @test size(gs.ndata.x) == size(X)
 
         gs = Zygote.gradient(model -> sum(node_features(model(g))), l)[1]
         @test size(gs.nn.weight) == size(l.nn.weight)