add neuralode example

Author: CarloLucibello

examples/Project.toml

@@ -1,8 +1,11 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DiffEqFlux = "aae7a2af-3d4f-5e19-a356-7da93b79d9d0"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"

examples/neural_ode.jl

@@ -0,0 +1,60 @@
+# Load the packages
+using GraphNeuralNetworks, JLD2, DiffEqFlux, DifferentialEquations
+using Flux: onehotbatch, onecold, throttle
+using Flux.Losses: logitcrossentropy
+using Statistics: mean
+using MLDatasets: Cora
+
+device = cpu
+
+# 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]
+
+
+# Model and Data Configuration
+nin = size(X,1)
+nhidden = 16
+nout = data.num_classes 
+epochs = 40
+
+# Define the Neural GDE
+diffeqarray_to_array(X) = reshape(cpu(X), size(X)[1:2])
+
+# GCNConv(nhidden => nhidden, graph=g),
+
+node = NeuralODE(
+    WithGraph(GCNConv(nhidden => nhidden), g),
+    (0.f0, 1.f0), Tsit5(), save_everystep = false,
+    reltol = 1e-3, abstol = 1e-3, save_start = false
+)
+
+model = GNNChain(GCNConv(nin => nhidden, relu),
+                Dropout(0.5),
+                node,
+                diffeqarray_to_array,
+                GCNConv(nhidden => nout))
+
+# Loss
+loss(x, y) = logitcrossentropy(model(g, x), y)
+accuracy(x, y) = mean(onecold(model(g, x)) .== onecold(y))
+
+# Training
+## Model Parameters
+ps = Flux.params(model, node.p);
+
+## Optimizer
+opt = ADAM(0.01)
+
+## Training Loop
+for epoch in 1:epochs
+    gs = gradient(() -> loss(X, y), ps)
+    Flux.Optimisers.update!(opt, ps, gs)
+    @show(accuracy(X, y))
+end