working tudataset example on cpu

Author: CarloLucibello

examples/graph_classification_mutag.jl renamed to examples/graph_classification_tudataset.jl

@@ -1,19 +1,28 @@
 # An example of semi-supervised node classification
 
 using Flux
-using Flux: @functor, dropout, onecold, onehotbatch
+using Flux: @functor, dropout, onecold, onehotbatch, getindex
 using Flux.Losses: logitbinarycrossentropy
+using Flux.Data: DataLoader
 using GraphNeuralNetworks
 using MLDatasets: TUDataset
 using Statistics, Random
 using CUDA
 CUDA.allowscalar(false)
 
-function eval_loss_accuracy(model, g, X, y)
-    ŷ = model(g, X) |> vec
-    l = logitbinarycrossentropy(ŷ, y)
-    acc = mean((2 .* ŷ .- 1) .* (2 .* y .- 1) .> 0)
-    return (loss = round(l, digits=4), acc = round(acc*100, digits=2))
+function eval_loss_accuracy(model, data_loader, device)
+    loss = 0.
+    acc = 0.
+    ntot = 0
+    for (g, X, y) in data_loader
+        g, X, y = g |> device, X |> device, y |> device
+        n = length(y) 
+        ŷ = model(g, X) |> vec
+        loss += logitbinarycrossentropy(ŷ, y) * n 
+        acc += mean((2 .* ŷ .- 1) .* (2 .* y .- 1) .> 0) * n
+        ntot += n
+    end        
+    return (loss = round(loss/ntot, digits=4), acc = round(acc*100/ntot, digits=2))
 end
 
 struct GNNData
@@ -22,6 +31,16 @@ struct GNNData
     y
 end
 
+Base.getindex(data::GNNData, i::Int) = getindex(data, [i])
+
+function Base.getindex(data::GNNData, i::AbstractVector)
+    sg, nodemap = subgraph(data.g, i)
+    return (sg, data.X[:,nodemap], data.y[i])
+end
+
+# Flux's Dataloader compatibility. 
+Flux.Data._nobs(data::GNNData) = data.g.num_graphs
+Flux.Data._getobs(data::GNNData, i) = data[i] 
 
 function getdataset(idxs)
     data = TUDataset("MUTAG")[idxs]
@@ -37,9 +56,10 @@ end
 # arguments for the `train` function 
 Base.@kwdef mutable struct Args
     η = 1f-3             # learning rate
-    epochs = 1000          # number of epochs
+    batchsize = 64      # batch size (number of graphs in each batch)
+    epochs = 200         # number of epochs
     seed = 17             # set seed > 0 for reproducibility
-    use_cuda = false      # if true use cuda (if available)
+    usecuda = true      # if true use cuda (if available)
     nhidden = 128        # dimension of hidden features
     infotime = 10 	     # report every `infotime` epochs
 end
@@ -48,7 +68,7 @@ function train(; kws...)
     args = Args(; kws...)
     args.seed > 0 && Random.seed!(args.seed)
 
-    if args.use_cuda && CUDA.functional()
+    if args.usecuda && CUDA.functional()
         device = gpu
         args.seed > 0 && CUDA.seed!(args.seed)
         @info "Training on GPU"
@@ -61,12 +81,15 @@ function train(; kws...)
 
     permindx = randperm(188)
     ntrain = 150
-    gtrain, Xtrain, ytrain = getdataset(permindx[1:ntrain]) 
-    gtest, Xtest, ytest = getdataset(permindx[ntrain+1:end]) 
+    dtrain = getdataset(permindx[1:ntrain]) 
+    dtest = getdataset(permindx[ntrain+1:end]) 
+    
+    train_loader = DataLoader(dtrain, batchsize=args.batchsize, shuffle=true)
+    test_loader = DataLoader(dtest, batchsize=args.batchsize, shuffle=false)
 
     # DEFINE MODEL
 
-    nin = size(Xtrain,1)
+    nin = size(dtrain.X, 1)
     nhidden = args.nhidden
 
     model = GNNChain(GCNConv(nin => nhidden, relu),
@@ -82,22 +105,23 @@ function train(; kws...)
     # LOGGING FUNCTION
 
     function report(epoch)
-        train = eval_loss_accuracy(model, gtrain, Xtrain, ytrain)
-        test = eval_loss_accuracy(model, gtest, Xtest, ytest)
+        train = eval_loss_accuracy(model, train_loader, device)
+        test = eval_loss_accuracy(model, test_loader, device)
         println("Epoch: $epoch   Train: $(train)   Test: $(test)")
     end
 
     # TRAIN
 
     report(0)
     for epoch in 1:args.epochs
-        # for (g, X, y) in train_loader
+        for (g, X, y) in train_loader
+            g, X, y = g |> device, X |> device, y |> device
             gs = Flux.gradient(ps) do
-                ŷ = model(gtrain, Xtrain) |> vec
-                logitbinarycrossentropy(ŷ, ytrain)
+                ŷ = model(g, X) |> vec
+                logitbinarycrossentropy(ŷ, y)
             end
             Flux.Optimise.update!(opt, ps, gs)
-        # end
+        end
 
         epoch % args.infotime == 0 && report(epoch)
     end

examples/node_classification_cora.jl

@@ -21,7 +21,7 @@ Base.@kwdef mutable struct Args
     η = 1f-3             # learning rate
     epochs = 100          # number of epochs
     seed = 17             # set seed > 0 for reproducibility
-    use_cuda = true      # if true use cuda (if available)
+    usecuda = true      # if true use cuda (if available)
     nhidden = 128        # dimension of hidden features
     infotime = 10 	     # report every `infotime` epochs
 end
@@ -33,7 +33,7 @@ function train(; kws...)
         CUDA.seed!(args.seed)
     end
 
-    if args.use_cuda && CUDA.functional()
+    if args.usecuda && CUDA.functional()
         device = gpu
         @info "Training on GPU"
     else