wip

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"

examples/cora.jl

@@ -18,8 +18,8 @@ end
 @functor GNN
 
 function GNN(; nin, nhidden, nout)
-    GNN(GCNConv(nin => nhidden, relu),
-        GCNConv(nhidden => nhidden, relu), 
+    GNN(GraphConv(nin => nhidden, relu),
+        GraphConv(nhidden => nhidden, relu), 
         Dense(nhidden, nout))
 end
 
@@ -35,13 +35,9 @@ 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 = l |> round4, acc = acc*100 |> round4)
+    return (loss = round(l, digits=4), acc = round(acc*100, digits=2))
 end
 
-## utility functions
-num_params(model) = sum(length, Flux.params(model)) 
-round4(x) = round(x, digits=4)
-
 # arguments for the `train` function 
 Base.@kwdef mutable struct Args
     η = 1f-3             # learning rate
@@ -68,7 +64,7 @@ function train(; kws...)
     end
 
     data = Cora.dataset()
-    fg = FeaturedGraph(data.adjacency_list)
+    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
@@ -85,9 +81,8 @@ function train(; kws...)
 
     function report(epoch)
         train = eval_loss_accuracy(X, y, train_ids, model, fg)
-        val = eval_loss_accuracy(X, y, val_ids, model, fg)
         test = eval_loss_accuracy(X, y, test_ids, model, fg)        
-        println("Epoch: $epoch   Train: $(train)  Val: $(val)   Test: $(test)")
+        println("Epoch: $epoch   Train: $(train)   Test: $(test)")
     end
 
     ## TRAINING