dataloader support; working gpu

Author: CarloLucibello

examples/Project.toml

@@ -4,3 +4,5 @@ Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 GraphNeuralNetworks = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"

examples/graph_classification_tudataset.jl

@@ -38,15 +38,14 @@ function Base.getindex(data::GNNData, i::AbstractVector)
     return (sg, data.X[:,nodemap], data.y[i])
 end
 
-# Flux's Dataloader compatibility. 
+# Flux's Dataloader compatibility. Related PR https://github.com/FluxML/Flux.jl/pull/1683
 Flux.Data._nobs(data::GNNData) = data.g.num_graphs
 Flux.Data._getobs(data::GNNData, i) = data[i] 
 
-function getdataset(idxs)
-    data = TUDataset("MUTAG")[idxs]
-    @info "MUTAG: num_nodes: $(data.num_nodes)  num_edges: $(data.num_edges)  num_graphs: $(data.num_graphs)"
+function process_dataset(data)
     g = GNNGraph(data.source, data.target, num_nodes=data.num_nodes, graph_indicator=data.graph_indicator)
     X = Array{Float32}(onehotbatch(data.node_labels, 0:6))
+    # The dataset also has edge features but we won't be using them
     # E = Array{Float32}(onehotbatch(data.edge_labels, sort(unique(data.edge_labels))))
     y = (1 .+ Array{Float32}(data.graph_labels)) ./ 2
     @assert all(∈([0,1]), y) # binary classification 
@@ -78,12 +77,18 @@ function train(; kws...)
     end
 
     # LOAD DATA
+
+    NUM_TRAIN = 150
+    full_data = TUDataset("MUTAG")
 
-    permindx = randperm(188)
-    ntrain = 150
-    dtrain = getdataset(permindx[1:ntrain]) 
-    dtest = getdataset(permindx[ntrain+1:end]) 
+    @info "MUTAG DATASET
+            num_nodes: $(full_data.num_nodes)  
+            num_edges: $(full_data.num_edges)  
+            num_graphs: $(full_data.num_graphs)"
 
+    perm = randperm(full_data.num_graphs)
+    dtrain = process_dataset(full_data[perm[1:NUM_TRAIN]]) 
+    dtest = process_dataset(full_data[perm[NUM_TRAIN+1:end]]) 
     train_loader = DataLoader(dtrain, batchsize=args.batchsize, shuffle=true)
     test_loader = DataLoader(dtest, batchsize=args.batchsize, shuffle=false)
 
@@ -92,9 +97,9 @@ function train(; kws...)
     nin = size(dtrain.X, 1)
     nhidden = args.nhidden
 
-    model = GNNChain(GCNConv(nin => nhidden, relu),
+    model = GNNChain(GraphConv(nin => nhidden, relu),
                      Dropout(0.5),
-                     GCNConv(nhidden => nhidden, relu),
+                     GraphConv(nhidden => nhidden, relu),
                      GlobalPool(mean), 
                      Dense(nhidden, 1))  |> device
 
@@ -127,4 +132,4 @@ function train(; kws...)
     end
 end
 
-# train()
+train()

src/gnngraph.jl

@@ -490,7 +490,6 @@ function subgraph(g::GNNGraph{<:COO_T}, i::AbstractVector)
     s = [nodemap[i] for i in s[edge_mask]]
     t = [nodemap[i] for i in t[edge_mask]]
     w = isnothing(w) ? nothing : w[edge_mask]
-    @show size(g.nf) size(node_mask)
     nf = isnothing(g.nf) ? nothing : g.nf[:,node_mask]
     ef = isnothing(g.ef) ? nothing : g.ef[:,edge_mask]
     gf = isnothing(g.gf) ? nothing : g.gf[:,i]

test/gnngraph.jl

@@ -17,7 +17,7 @@
         @test sort(outneighbors(g, 1)) == [2, 4] 
         @test sort(inneighbors(g, 1)) == [2, 4] 
         @test is_directed(g) == true
-        s1, t1 = sort_edge_index(edge_index(g))
+        s1, t1 = GraphNeuralNetworks.sort_edge_index(edge_index(g))
         @test s1 == s
         @test t1 == t
 
@@ -65,7 +65,7 @@
         @test sort(outneighbors(g, 1)) == [2] 
         @test sort(inneighbors(g, 1)) == [4] 
         @test is_directed(g) == true
-        s1, t1 = sort_edge_index(edge_index(g))
+        s1, t1 = GraphNeuralNetworks.sort_edge_index(edge_index(g))
         @test s1 == s
         @test t1 == t
 
@@ -103,6 +103,7 @@
     end
 
     @testset "batch"  begin
+        #TODO add graph_type=GRAPH_T
         g1 = GNNGraph(random_regular_graph(10,2), nf=rand(16,10))
         g2 = GNNGraph(random_regular_graph(4,2), nf=rand(16,4))
         g3 = GNNGraph(random_regular_graph(7,2), nf=rand(16,7))
@@ -112,5 +113,25 @@
 
         g123 = Flux.batch([g1, g2, g3])
         @test g123.graph_indicator == [fill(1, 10); fill(2, 4); fill(3, 7)]
+
+        s, t = edge_index(g123)
+        @test s == [edge_index(g1)[1]; 10 .+ edge_index(g2)[1]; 14 .+ edge_index(g3)[1]] 
+        @test t == [edge_index(g1)[2]; 10 .+ edge_index(g2)[2]; 14 .+ edge_index(g3)[2]] 
+        @test g123.nf[:,11:14] ≈ g2.nf 
+    end
+
+    @testset "subgraph"  begin
+        #TODO add graph_type=GRAPH_T
+        g1 = GNNGraph(random_regular_graph(10,2), nf=rand(16,10))
+        g2 = GNNGraph(random_regular_graph(4,2), nf=rand(16,4))
+        g3 = GNNGraph(random_regular_graph(7,2), nf=rand(16,7))
+        g = Flux.batch([g1, g2, g3])
+        g2b, nodemap = subgraph(g, 2)
+        
+        s, t = edge_index(g2b)
+        @test s == edge_index(g2)[1]
+        @test t == edge_index(g2)[2] 
+        @test g2b.nf ≈ g2.nf 
     end
+
 end

test/runtests.jl

@@ -1,5 +1,4 @@
 using GraphNeuralNetworks
-using GraphNeuralNetworks: sort_edge_index
 using Flux
 using CUDA
 using Flux: gpu, @functor