neural ode example working on cpu and gpu

Author: CarloLucibello

examples/neural_ode.jl

@@ -1,11 +1,12 @@
 # Load the packages
-using GraphNeuralNetworks, JLD2, DiffEqFlux, DifferentialEquations
-using Flux: onehotbatch, onecold, throttle
+using GraphNeuralNetworks, DiffEqFlux, DifferentialEquations
+using Flux: onehotbatch, onecold
 using Flux.Losses: logitcrossentropy
 using Statistics: mean
 using MLDatasets: Cora
 
-device = cpu # `gpu` not working yet
+# device = cpu # `gpu` not working yet
+device = gpu
 
 # LOAD DATA
 data = Cora.dataset()
@@ -39,21 +40,21 @@ node = NeuralODE(WithGraph(node_chain, g),
 model = GNNChain(GCNConv(nin => nhidden, relu),
                  Dropout(0.5),
                  node,
-                 diffeqarray_to_array,
+                 diffeqsol_to_array,
                  Dense(nhidden, nout)) |> device
 
 # 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);
+# # Training
+# ## Model Parameters
+ps = Flux.params(model);
 
-## Optimizer
+# ## Optimizer
 opt = ADAM(0.01)
 
-## Training Loop
+# ## Training Loop
 for epoch in 1:epochs
     gs = gradient(() -> loss(X, y), ps)
     Flux.Optimise.update!(opt, ps, gs)

src/GNNGraphs/gnngraph.jl

@@ -5,8 +5,8 @@ https://juliagraphs.org/Graphs.jl/latest/types/#AbstractGraph-Type
 https://juliagraphs.org/Graphs.jl/latest/developing/#Developing-Alternate-Graph-Types
 =============================================#
 
-const COO_T = Tuple{T, T, V} where {T <: AbstractVector, V}
-const ADJLIST_T = AbstractVector{T} where T <: AbstractVector
+const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V}
+const ADJLIST_T = AbstractVector{T} where T <: AbstractVector{<:Integer}
 const ADJMAT_T = AbstractMatrix
 const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T
 const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}

src/layers/basic.jl

@@ -13,7 +13,7 @@ abstract type GNNLayer end
 
 
 """
-  WithGraph(model, g::GNNGraph; traingraph=false) 
+    WithGraph(model, g::GNNGraph; traingraph=false) 
 
 A type wrapping the `model` and tying it to the graph `g`.
 In the forward pass, can only take feature arrays as inputs,
@@ -38,17 +38,31 @@ x2 = rand(Float32, 2, 4)
 @assert wg(g2, x2) == model(g2, x2)
 ```
 """
-struct WithGraph{M}
-  model::M
-  g::GNNGraph
-  traingraph::Bool
+struct WithGraph{M, G<:GNNGraph}
+    model::M
+    g::G
+    traingraph::Bool
 end
 
 WithGraph(model, g::GNNGraph; traingraph=false) = WithGraph(model, g, traingraph)
 
 @functor WithGraph
 Flux.trainable(l::WithGraph) = l.traingraph ? (l.model, l.g) : (l.model,)
 
+# Work around 
+# https://github.com/FluxML/Flux.jl/issues/1733
+# Revisit after 
+# https://github.com/FluxML/Flux.jl/pull/1742
+function Flux.destructure(m::WithGraph)
+    @assert m.traingraph == false # TODO
+    p, re = Flux.destructure(m.model)
+    function  re_withgraph(x)
+        WithGraph(re(x), m.g, m.traingraph)        
+    end
+
+    return p, re_withgraph
+end
+
 (l::WithGraph)(g::GNNGraph, x...; kws...) = l.model(g, x...; kws...)
 (l::WithGraph)(x...; kws...) = l.model(l.g, x...; kws...)
 
@@ -86,15 +100,15 @@ julia> m(g, x)
 ```
 """
 struct GNNChain{T} <: GNNLayer
-  layers::T
-
-  GNNChain(xs...) = new{typeof(xs)}(xs)
-  
-  function GNNChain(; kw...)
-    :layers in Base.keys(kw) && throw(ArgumentError("a GNNChain cannot have a named layer called `layers`"))
-    isempty(kw) && return new{Tuple{}}(())
-    new{typeof(values(kw))}(values(kw))
-  end
+    layers::T
+
+    GNNChain(xs...) = new{typeof(xs)}(xs)
+    
+    function GNNChain(; kw...)
+        :layers in Base.keys(kw) && throw(ArgumentError("a GNNChain cannot have a named layer called `layers`"))
+        isempty(kw) && return new{Tuple{}}(())
+        new{typeof(values(kw))}(values(kw))
+    end
 end
 
 @forward GNNChain.layers Base.getindex, Base.length, Base.first, Base.last,

test/examples/node_classification_cora.jl

@@ -14,7 +14,6 @@ function eval_loss_accuracy(X, y, ids, model, g)
     return (loss = round(l, digits=4), acc = round(acc*100, digits=2))
 end
 
-
 # arguments for the `train` function 
 Base.@kwdef mutable struct Args
     η = 5f-3             # learning rate