Merge pull request #9 from CarloLucibello/cl/cora

add Cora example

Author: CarloLucibello

Project.toml

@@ -11,6 +11,7 @@ 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"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

README.md

@@ -21,60 +21,7 @@ Some of its noticeable features are the following:
 ]add GraphNeuralNetworks
 ```
 
-## Featured Graphs
+## Usage
 
-GraphNeuralNetworks handles graph data (the graph topology + node/edge/global features)
-thanks to the type `FeaturedGraph`.
-
-A `FeaturedGraph` can be constructed out of 
-adjacency matrices, adjacency lists, LightGraphs' types...
-
-```julia
-fg = FeaturedGraph(adj_list)   
-```
-
-## Graph convolutional layers
-
-Construct a GCN layer:
-
-```julia
-GCNConv(input_dim => output_dim, relu)
-```
-
-## Usage Example
-
-```julia
-struct GNN
-    conv1
-    conv2 
-    dense
-end
-
-@functor GNN
-
-function GNN()
-    GNN(GCNConv(1024=>512, relu),
-        GCNConv(512=>128, relu), 
-        Dense(128, 10))
-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
-
-model = GNN()
-
-loss(x, y) = logitcrossentropy(model(fg, x), y)
-accuracy(x, y) = mean(onecold(model(fg, x)) .== onecold(y))
-
-ps = Flux.params(model)
-train_data = [(train_X, train_y)]
-opt = ADAM(0.01)
-evalcb() = @show(accuracy(train_X, train_y))
-
-Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))
+Usage examples can be found in the `examples/` folder. 
 ```

examples/cora.jl

@@ -0,0 +1,103 @@
+# An example of semi-supervised node classification
+
+using Flux
+using Flux: @functor, dropout, onecold, onehotbatch
+using Flux.Losses: logitcrossentropy
+using GraphNeuralNetworks
+using MLDatasets: Cora
+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)
+    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))
+end
+
+# arguments for the `train` function 
+Base.@kwdef mutable struct Args
+    η = 1f-3             # learning rate
+    epochs = 100          # number of epochs
+    seed = 17             # set seed > 0 for reproducibility
+    use_cuda = true      # if true use cuda (if available)
+    nhidden = 128        # dimension of hidden features
+    infotime = 10 	     # report every `infotime` epochs
+end
+
+function train(; kws...)
+    args = Args(; kws...)
+    if args.seed > 0
+        Random.seed!(args.seed)
+        CUDA.seed!(args.seed)
+    end
+    
+    if args.use_cuda && CUDA.functional()
+        device = gpu
+        @info "Training on GPU"
+    else
+        device = cpu
+        @info "Training on CPU"
+    end
+
+    data = Cora.dataset()
+    fg = FeaturedGraph(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
+    ps = Flux.params(model)
+    opt = ADAM(args.η)
+
+    @info "NUM NODES: $(fg.num_nodes)  NUM EDGES: $(fg.num_edges)"
+    
+    function report(epoch)
+        train = eval_loss_accuracy(X, y, train_ids, model, fg)
+        test = eval_loss_accuracy(X, y, test_ids, model, fg)        
+        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)
+            logitcrossentropy(ŷ[:,train_ids], ytrain)
+        end
+
+        Flux.Optimise.update!(opt, ps, gs)
+        
+        epoch % args.infotime == 0 && report(epoch)
+    end
+end
+
+train()

src/featuredgraph.jl

@@ -234,10 +234,10 @@ function LightGraphs.adjacency_matrix(fg::FeaturedGraph{<:ADJMAT_T}, T::DataType
     return dir == :out ? A : A'
 end
 
-function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}; dir=:out)
+function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}, T=Int; dir=:out)
     s, t = edge_index(fg)
-    degs = fill!(similar(s, eltype(s), fg.num_nodes), 0)
-    o = fill!(similar(s, eltype(s), fg.num_edges), 1)
+    degs = fill!(similar(s, T, fg.num_nodes), 0)
+    o = fill!(similar(s, Int, fg.num_edges), 1)
     if dir ∈ [:out, :both]
         NNlib.scatter!(+, degs, o, s)
     end
@@ -247,9 +247,9 @@ function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}; dir=:out)
     return degs
 end
 
-function LightGraphs.degree(fg::FeaturedGraph{<:ADJMAT_T}; dir=:out)
+function LightGraphs.degree(fg::FeaturedGraph{<:ADJMAT_T}, T=Int; dir=:out)
     @assert dir ∈ (:in, :out)
-    A = graph(fg)
+    A = adjacency_matrix(fg, T)
     return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
 end
 
@@ -298,29 +298,32 @@ function LightGraphs.laplacian_matrix(fg::FeaturedGraph, T::DataType=Int; dir::S
 end
 
 """
-    normalized_laplacian(fg, T=Float32; selfloop=false, dir=:out)
+    normalized_laplacian(fg, T=Float32; add_self_loops=false, dir=:out)
 
 Normalized Laplacian matrix of graph `g`.
 
 # Arguments
 
 - `fg`: A `FeaturedGraph`.
 - `T`: result element type.
-- `selfloop`: adding self loop while calculating the matrix.
+- `add_self_loops`: add self-loops while calculating the matrix.
 - `dir`: the edge directionality considered (:out, :in, :both).
 """
-function normalized_laplacian(fg::FeaturedGraph, T::DataType=Float32; selfloop::Bool=false, dir::Symbol=:out)
+function normalized_laplacian(fg::FeaturedGraph, T::DataType=Float32; 
+                        add_self_loops::Bool=false, dir::Symbol=:out)
+    Ã = normalized_adjacency(fg, T; dir, add_self_loops)
+    return I - Ã
+end
+
+function normalized_adjacency(fg::FeaturedGraph, T::DataType=Float32; 
+                        add_self_loops::Bool=false, dir::Symbol=:out)
     A = adjacency_matrix(fg, T; dir=dir)
-    sz = size(A)
-    @assert sz[1] == sz[2]
-    if selfloop
-        A += I - Diagonal(A)
-    else
-        A -= Diagonal(A) 
+    if add_self_loops
+        A += I
     end
     degs = vec(sum(A; dims=2))
     inv_sqrtD = Diagonal(inv.(sqrt.(degs)))
-    return I - inv_sqrtD * A * inv_sqrtD
+    return inv_sqrtD * A * inv_sqrtD
 end
 
 @doc raw"""
@@ -354,14 +357,14 @@ _eigmax(A) = KrylovKit.eigsolve(Symmetric(A), 1, :LR)[1][1] # also eigs(A, x0, n
 
 Return a featured graph with the same features as `fg`
 but also adding edges connecting the nodes to themselves.
+
+Nodes with already existing
+self-loops will obtain a second self-loop.
 """
 function add_self_loops(fg::FeaturedGraph{<:COO_T})
     s, t = edge_index(fg)
     @assert edge_feature(fg) === nothing
     @assert edge_weight(fg) === nothing
-    mask_old_loops = s .!= t
-    s = s[mask_old_loops]
-    t = t[mask_old_loops]
     n = fg.num_nodes
     nodes = convert(typeof(s), [1:n;])
     s = [s; nodes]
@@ -371,12 +374,11 @@ function add_self_loops(fg::FeaturedGraph{<:COO_T})
         node_feature(fg), edge_feature(fg), global_feature(fg))
 end
 
-function add_self_loops(fg::FeaturedGraph{<:ADJMAT_T})
+function add_self_loops(fg::FeaturedGraph{<:ADJMAT_T}; add_to_existing=true)
     A = graph(fg)
     @assert edge_feature(fg) === nothing
-    nold = sum(Diagonal(A)) |> Int
-    A = A - Diagonal(A) + I
-    num_edges =  fg.num_edges - nold + fg.num_nodes
+    A += I
+    num_edges =  fg.num_edges + fg.num_nodes
     FeaturedGraph(A, fg.num_nodes, num_edges,
         node_feature(fg), edge_feature(fg), global_feature(fg))
 end
@@ -396,6 +398,7 @@ function remove_self_loops(fg::FeaturedGraph{<:COO_T})
 end
 
 @non_differentiable normalized_laplacian(x...)
+@non_differentiable normalized_adjacency(x...)
 @non_differentiable scaled_laplacian(x...)
 @non_differentiable adjacency_matrix(x...)
 @non_differentiable adjacency_list(x...)

src/layers/conv.jl

@@ -34,22 +34,20 @@ end
 ## but cannot compute the normalized laplacian of sparse cuda matrices yet,
 ## therefore fallback to message passing framework on gpu for the time being
 
-function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
-    L̃ = normalized_laplacian(fg, eltype(x); selfloop=true)
-    l.σ.(l.weight * x * L̃ .+ l.bias)
+function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix{T}) where T
+    Ã = normalized_adjacency(fg, T; dir=:out, add_self_loops=true)
+    l.σ.(l.weight * x * Ã .+ l.bias)
 end
 
 message(l::GCNConv, xi, xj) = xj
 update(l::GCNConv, m, x) = m
 
-function (l::GCNConv)(fg::FeaturedGraph, x::CuMatrix)
+function (l::GCNConv)(fg::FeaturedGraph, x::CuMatrix{T}) where T
     fg = add_self_loops(fg)
-    T = eltype(l.weight)
-    # cout = sqrt.(degree(fg, dir=:out))
-    cin = 1 ./ reshape(sqrt.(T.(degree(fg, dir=:in))), 1, :)
-    x = cin .* x
+    c = 1 ./ sqrt.(degree(fg, T, dir=:in))
+    x = x .* c'
     _, x = propagate(l, fg, nothing, x, nothing, +)
-    x = cin .* x
+    x = x .* c'
     return l.σ.(l.weight * x .+ l.bias)
 end