fix examples

fix GAE, VGAE, GDEs

Author: yuehhua

examples/gae.jl

@@ -1,4 +1,5 @@
 using GeometricFlux
+using GraphSignals
 using Flux
 using Flux: throttle
 using Flux.Losses: logitbinarycrossentropy
@@ -9,6 +10,8 @@ using SparseArrays
 using Graphs.SimpleGraphs
 using CUDA
 
+CUDA.allowscalar(false)
+
 @load "data/cora_features.jld2" features
 @load "data/cora_graph.jld2" g
 
@@ -20,14 +23,15 @@ target_catg = 7
 epochs = 200
 
 ## Preprocessing data
-fg = FeaturedGraph(g) |> gpu
+fg = FeaturedGraph(g)  # pass to gpu together in model layers
 train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
-train_y = fg  # dim: num_nodes * num_nodes
+train_y = fg |> GraphSignals.adjacency_matrix |> gpu  # dim: num_nodes * num_nodes
 
 ## Model
 encoder = Chain(GCNConv(fg, num_features=>hidden1, relu),
                 GCNConv(fg, hidden1=>hidden2))
-model = Chain(GAE(encoder, σ)) |> gpu
+model = Chain(GAE(encoder, σ)) |> gpu;
+# do not show model architecture, showing CuSparseMatrix will trigger errors
 
 ## Loss
 loss(x, y) = logitbinarycrossentropy(model(x), y)

examples/gat.jl

@@ -2,10 +2,10 @@ using GeometricFlux
 using Flux
 using Flux: onehotbatch, onecold, logitcrossentropy, throttle
 using Flux: @epochs
-using Flux.Data: DataLoader
 using JLD2
 using Statistics: mean
 using SparseArrays
+using LinearAlgebra
 using Graphs.SimpleGraphs
 using Graphs: adjacency_matrix
 using CUDA
@@ -24,29 +24,30 @@ epochs = 10
 ## Preprocessing data
 train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
 train_y = Matrix{Float32}(labels) |> gpu  # dim: target_catg * num_nodes
-adj_mat = Matrix{Float32}(adjacency_matrix(g)) |> gpu
+A = Matrix{Int}((adjacency_matrix(g) + I) .≥ 1)
+fg = FeaturedGraph(A, :adjm)
 
 ## Model
-model = Chain(GATConv(g, num_features=>hidden, heads=heads),
+model = Chain(GATConv(fg, num_features=>hidden, heads=heads),
               Dropout(0.6),
-              GATConv(g, hidden*heads=>target_catg, heads=heads, concat=false)
+              GATConv(fg, hidden*heads=>target_catg, heads=heads, concat=false)
               ) |> gpu
 # test model
-# @show model(train_X)
+@show model(train_X)
 
 ## Loss
 loss(x, y) = logitcrossentropy(model(x), y)
 accuracy(x, y) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))
 
 # test loss
-# @show loss(train_X, train_y)
+@show loss(train_X, train_y)
 
 # test gradient
-# @show gradient(X -> loss(X, train_y), train_X)
+@show gradient(()->loss(train_X, train_y), Flux.params(model))
 
 ## Training
 ps = Flux.params(model)
-train_data = DataLoader(train_X, train_y, batchsize=num_nodes)
+train_data = Flux.Data.DataLoader((train_X, train_y), batchsize=num_nodes)
 opt = ADAM(0.01)
 evalcb() = @show(accuracy(train_X, train_y))
 

examples/gde.jl

@@ -1,4 +1,4 @@
-using GeometricFlux, Flux, JLD2, SparseArrays, DiffEqFlux, DifferentialEquations
+using GeometricFlux, GraphSignals, Flux, JLD2, SparseArrays, DiffEqFlux, DifferentialEquations
 using Flux: onehotbatch, onecold, logitcrossentropy, throttle
 using Flux: @epochs
 using Statistics: mean

examples/gde_gpu.jl

@@ -1,4 +1,4 @@
-using GeometricFlux, Flux, JLD2, SparseArrays, DiffEqFlux, DifferentialEquations
+using GeometricFlux, GraphSignals, Flux, JLD2, SparseArrays, DiffEqFlux, DifferentialEquations
 using Flux: onehotbatch, onecold, logitcrossentropy, throttle
 using Flux: @epochs
 using Statistics: mean
@@ -19,7 +19,7 @@ epochs = 40
 # Preprocess the data and compute adjacency matrix
 train_X = Matrix{Float32}(features) |> gpu  # dim: num_features * num_nodes
 train_y = Float32.(labels) |> gpu  # dim: target_catg * num_nodes
-fg = FeaturedGraph(g) |> gpu
+fg = FeaturedGraph(g)
 
 # Define the Neural GDE
 diffeqarray_to_array(x) = reshape(gpu(x), size(x)[1:2])

examples/vgae.jl

@@ -23,7 +23,7 @@ epochs = 200
 ## Preprocessing data
 masks = [rand(Float32, num_nodes, num_nodes).>0.1 for i in 1:10]
 adj_mat = Matrix{Float32}(adjacency_matrix(g))
-train_data = [(FeaturedGraph(adj_mat.*M, Matrix{Float32}(features)), adj_mat) for M in masks]
+train_data = [(FeaturedGraph(adj_mat.*M, nf=Matrix{Float32}(features)), adj_mat) for M in masks]
 
 ## Model
 model = VGAE(GCNConv(num_features=>h_dim, relu;), h_dim, z_dim, σ)
@@ -34,7 +34,7 @@ ps = Flux.params(model)
 l2_norm(p) = sum(abs2, p)
 
 function loss(fg, Y, X=node_feature(fg), T=eltype(X), β=one(T), λ=T(0.01); debug=false)
-    μ̂, logσ̂ = summarize(encoder, fg)
+    μ̂, logσ̂ = GeometricFlux.summarize(encoder, fg)
     Z = node_feature(encoder(fg))
     kl_q_p = -T(0.5) * sum(one(T) .+ T(2).*logσ̂ .- μ̂.^2 .- exp.(T(2).*logσ̂))
     logp_y_z = -sum(logitbinarycrossentropy(decoder(Z), Y, agg=identity)) / size(Y,2)