fix gpu tests

Author: CarloLucibello

GraphNeuralNetworks/test/examples/node_classification_cora.jl

@@ -1,14 +1,11 @@
-@testitem "Training Example" setup=[TestModule] begin
-    using .TestModule
+@testmodule TrainingExampleModule begin
     using Flux
     using Flux: onecold, onehotbatch
     using Flux.Losses: logitcrossentropy
     using GraphNeuralNetworks
     using MLDatasets: Cora
     using Statistics, Random
-    using CUDA
-    CUDA.allowscalar(false)
-
+    
     function eval_loss_accuracy(X, y, ids, model, g)
         ŷ = model(g, X)
         l = logitcrossentropy(ŷ[:, ids], y[:, ids])
@@ -21,27 +18,27 @@
         η = 5.0f-3            # learning rate
         epochs = 10         # number of epochs
         seed = 17           # set seed > 0 for reproducibility
-        usecuda = false     # if true use cuda (if available)
+        use_gpu = false     # if true use gpu (if available)
         nhidden = 64        # dimension of hidden features
     end
 
     function train(Layer; verbose = false, kws...)
         args = Args(; kws...)
         args.seed > 0 && Random.seed!(args.seed)
 
-        if args.usecuda && CUDA.functional()
-            device = Flux.gpu
-            args.seed > 0 && CUDA.seed!(args.seed)
+        if args.use_gpu
+            device = gpu_device(force=true)
+            Random.seed!(default_device_rng(device))
         else
-            device = Flux.cpu
+            device = cpu_device()
         end
 
         # LOAD DATA
         dataset = Cora()
         classes = dataset.metadata["classes"]
         g = mldataset2gnngraph(dataset) |> device
         X = g.ndata.features
-        y = onehotbatch(g.ndata.targets |> cpu, classes) |> device # remove when https://github.com/FluxML/Flux.jl/pull/1959 tagged
+        y = onehotbatch(g.ndata.targets, classes)
         train_mask = g.ndata.train_mask
         test_mask = g.ndata.test_mask
         ytrain = y[:, train_mask]
@@ -78,7 +75,7 @@
         return train_res, test_res
     end
 
-    function train_many(; usecuda = false)
+    function train_many(; use_gpu = false)
         for (layer, Layer) in [
             ("GCNConv", (nin, nout) -> GCNConv(nin => nout, relu)),
             ("ResGatedGraphConv", (nin, nout) -> ResGatedGraphConv(nin => nout, relu)),
@@ -96,16 +93,21 @@
             ## ("EdgeConv",(nin, nout) -> EdgeConv(Dense(2nin, nout, relu))), # Fits the training set but does not generalize well
         ]
             @show layer
-            @time train_res, test_res = train(Layer; usecuda, verbose = false)
+            @time train_res, test_res = train(Layer; use_gpu, verbose = false)
             # @show train_res, test_res
             @test train_res.acc > 94
             @test test_res.acc > 69
         end
     end
+end # module
+
+@testitem "training example" setup=[TrainingExampleModule] begin
+    using .TrainingExampleModule
+    train_many()
+end
 
-    train_many(usecuda = false)
-    # #TODO
-    # if TEST_GPU
-    #     train_many(usecuda = true)
-    # end
+@testitem "training example GPU" setup=[TrainingExampleModule] begin
+    using .TrainingExampleModule
+    train_many(use_gpu = true)
 end
+

GraphNeuralNetworks/test/layers/conv.jl

@@ -252,6 +252,16 @@ end
     end
 end
 
+@testitem "EdgeConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
+    using .TestModule
+    l = EdgeConv(Dense(2 * D_IN, D_OUT), aggr = +)
+    for g in TEST_GRAPHS
+        g.graph isa AbstractSparseMatrix && continue
+        @test size(l(g, g.x)) == (D_OUT, g.num_nodes)
+        test_gradients(l, g, g.x, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
+    end   
+end
+
 @testitem "GINConv" setup=[TolSnippet, TestModule] begin
     using .TestModule
     nn = Dense(D_IN, D_OUT)
@@ -289,17 +299,17 @@ end
     end
 end
 
-@testitem "NNConv" setup=[TolSnippet, TestModule] begin
+@testitem "NNConv GPU" setup=[TolSnippet, TestModule] tags=[:gpu] begin
     using .TestModule
     edim = 10
     nn = Dense(edim, D_OUT * D_IN)
-
     l = NNConv(D_IN => D_OUT, nn, tanh, bias = true, aggr = +)
     for g in TEST_GRAPHS
+        g.graph isa AbstractSparseMatrix && continue
         g = GNNGraph(g, edata = rand(Float32, edim, g.num_edges))
         @test size(l(g, g.x, g.e)) == (D_OUT, g.num_nodes)
-        test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH)
-    end
+        test_gradients(l, g, g.x, g.e, rtol = RTOL_HIGH, test_gpu = true, compare_finite_diff = false)
+    end   
 end
 
 @testitem "SAGEConv" setup=[TolSnippet, TestModule] begin

GraphNeuralNetworks/test/test_module.jl

@@ -11,7 +11,7 @@ using Pkg
 
 ## Uncomment below to change the default test settings
 # ENV["GNN_TEST_CPU"] = "false"
-# ENV["GNN_TEST_CUDA"] = "true"
+ENV["GNN_TEST_CUDA"] = "true"
 # ENV["GNN_TEST_AMDGPU"] = "true"
 # ENV["GNN_TEST_Metal"] = "true"
 
@@ -24,7 +24,7 @@ if get(ENV, "GNN_TEST_AMDGPU", "false") == "true"
     Pkg.add("AMDGPU")
     using AMDGPU
 end
-if get(ENV, "GNN_TEST_Metal", "true") == "true"
+if get(ENV, "GNN_TEST_Metal", "false") == "true"
     Pkg.add("Metal")
     using Metal
 end