test examples

Author: CarloLucibello

Project.toml

@@ -12,6 +12,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"
@@ -36,8 +37,9 @@ julia = "1.6"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Adapt", "Zygote", "FiniteDifferences", "ChainRulesTestUtils"]
+test = ["Test", "Adapt", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "MLDatasets"]

examples/graph_classification_tudataset.jl

@@ -1,7 +1,7 @@
 # An example of graph classification
 
 using Flux
-using Flux: @functor, dropout, onecold, onehotbatch, getindex
+using Flux: @functor, dropout, onecold, onehotbatch, getindex, cpu, gpu
 using Flux.Losses: logitbinarycrossentropy
 using Flux.Data: DataLoader
 using GraphNeuralNetworks

examples/node_classification_cora.jl

@@ -1,7 +1,7 @@
 # An example of semi-supervised node classification
 
 using Flux
-using Flux: @functor, dropout, onecold, onehotbatch
+using Flux: onecold, onehotbatch
 using Flux.Losses: logitcrossentropy
 using GraphNeuralNetworks
 using MLDatasets: Cora
@@ -28,13 +28,12 @@ end
 
 function train(; kws...)
     args = Args(; kws...)
-    if args.seed > 0
-        Random.seed!(args.seed)
-        CUDA.seed!(args.seed)
-    end
+
+    args.seed > 0 && Random.seed!(args.seed)
 
     if args.usecuda && CUDA.functional()
         device = gpu
+        args.seed > 0 && CUDA.seed!(args.seed)
         @info "Training on GPU"
     else
         device = cpu

src/layers/conv.jl

@@ -288,9 +288,8 @@ end
 
 
 function Base.show(io::IO, l::GATConv)
-    in_channel = size(l.weight, ndims(l.weight))
-    out_channel = size(l.weight, ndims(l.weight)-1)
-    print(io, "GATConv(", in_channel, "=>", out_channel)
+    out_channel, in_channel = size(l.weight)
+    print(io, "GATConv(", in_channel, "=>", out_channel ÷ l.heads)
     print(io, ", LeakyReLU(λ=", l.negative_slope)
     print(io, "))")
 end
@@ -567,9 +566,8 @@ function (l::SAGEConv)(g::GNNGraph, x::AbstractMatrix)
 end
 
 function Base.show(io::IO, l::SAGEConv)
-    in_channel = size(l.weight1, ndims(l.weight1))
-    out_channel = size(l.weight1, ndims(l.weight1)-1)
-    print(io, "SAGEConv(", in_channel, " => ", out_channel)
+    out_channel, in_channel = size(l.weight)
+    print(io, "SAGEConv(", in_channel ÷ 2, " => ", out_channel)
     l.σ == identity || print(io, ", ", l.σ)
     print(io, ", aggr=", l.aggr)
     print(io, ")")

test/examples/node_classification_cora.jl

@@ -0,0 +1,98 @@
+using Flux
+using Flux: onecold, onehotbatch
+using Flux.Losses: logitcrossentropy
+using GraphNeuralNetworks
+using MLDatasets: Cora
+using Statistics, Random
+using CUDA
+CUDA.allowscalar(false)
+
+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))
+end
+
+
+# arguments for the `train` function 
+Base.@kwdef mutable struct Args
+    η = 1f-3             # learning rate
+    epochs = 20         # number of epochs
+    seed = 17             # set seed > 0 for reproducibility
+    usecuda = false      # if true use cuda (if available)
+    nhidden = 128        # dimension of hidden features
+end
+
+function train(Layer; verbose=false, kws...)
+    args = Args(; kws...)
+    args.seed > 0 && Random.seed!(args.seed)
+    
+    if args.usecuda && CUDA.functional()
+        device = Flux.gpu
+        args.seed > 0 && CUDA.seed!(args.seed)
+    else
+        device = Flux.cpu
+    end
+
+    # LOAD DATA
+    data = Cora.dataset()
+    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]
+
+    nin, nhidden, nout = size(X,1), args.nhidden, data.num_classes 
+    
+    ## DEFINE MODEL
+    model = GNNChain(Layer(nin, nhidden),
+                     Dropout(0.5),
+                     Layer(nhidden, nhidden), 
+                     Dense(nhidden, nout))  |> device
+
+    ps = Flux.params(model)
+    opt = ADAM(args.η)
+    
+
+    ## TRAINING
+    function report(epoch)
+        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
+    
+    verbose && report(0)
+    for epoch in 1:args.epochs
+        gs = Flux.gradient(ps) do
+            ŷ = model(g, X)
+            logitcrossentropy(ŷ[:,train_ids], ytrain)
+        end
+        verbose && report(epoch)
+        Flux.Optimise.update!(opt, ps, gs)
+    end
+
+    train_res = eval_loss_accuracy(X, y, train_ids, model, g)
+    test_res = eval_loss_accuracy(X, y, test_ids, model, g)        
+    return train_res, test_res
+end
+
+for Layer in [
+            (nin, nout) -> GCNConv(nin => nout, relu),
+            (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) -> 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=true)
+    # @show Layer(2,2) train_res, test_res
+    @test train_res.acc > 95
+    @test test_res.acc > 70
+end

test/layers/conv.jl

@@ -40,7 +40,7 @@
     end
 
     @testset "ChebConv" begin
-        k = 6
+        k = 3
         l = ChebConv(in_channel => out_channel, k)
         @test size(l.weight) == (out_channel, in_channel, k)
         @test size(l.bias) == (out_channel,)

test/runtests.jl

@@ -18,6 +18,7 @@ tests = [
     "layers/basic",
     "layers/conv",
     "layers/pool",
+    "examples.jl",
 ]
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")