improve gpu tests and some fixes

Author: CarloLucibello

src/featuredgraph.jl

@@ -234,10 +234,10 @@ function LightGraphs.adjacency_matrix(fg::FeaturedGraph{<:ADJMAT_T}, T::DataType
     return dir == :out ? A : A'
 end
 
-function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}; dir=:out)
+function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}, T=Int; dir=:out)
     s, t = edge_index(fg)
-    degs = fill!(similar(s, eltype(s), fg.num_nodes), 0)
-    o = fill!(similar(s, eltype(s), fg.num_edges), 1)
+    degs = fill!(similar(s, T, fg.num_nodes), 0)
+    o = fill!(similar(s, Int, fg.num_edges), 1)
     if dir ∈ [:out, :both]
         NNlib.scatter!(+, degs, o, s)
     end
@@ -247,9 +247,9 @@ function LightGraphs.degree(fg::FeaturedGraph{<:COO_T}; dir=:out)
     return degs
 end
 
-function LightGraphs.degree(fg::FeaturedGraph{<:ADJMAT_T}; dir=:out)
+function LightGraphs.degree(fg::FeaturedGraph{<:ADJMAT_T}, T=Int; dir=:out)
     @assert dir ∈ (:in, :out)
-    A = graph(fg)
+    A = adjacency_matrix(fg, T)
     return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
 end
 

src/layers/conv.jl

@@ -34,22 +34,20 @@ end
 ## but cannot compute the normalized laplacian of sparse cuda matrices yet,
 ## therefore fallback to message passing framework on gpu for the time being
 
-function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
-    Ã = normalized_adjacency(fg, eltype(x); dir=:out, add_self_loops=true)
-    l.σ.(l.weight * x * Ã .+ l.bias)
-end
+# function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
+#     Ã = normalized_adjacency(fg, eltype(x); dir=:out, add_self_loops=true)
+#     l.σ.(l.weight * x * Ã .+ l.bias)
+# end
 
 message(l::GCNConv, xi, xj) = xj
 update(l::GCNConv, m, x) = m
 
-function (l::GCNConv)(fg::FeaturedGraph, x::CuMatrix)
+function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix{T}) where T
     fg = add_self_loops(fg; add_to_existing=true)
-    T = eltype(l.weight)
-    # cout = sqrt.(degree(fg, dir=:out))
-    cin = 1 ./ reshape(sqrt.(T.(degree(fg, dir=:in))), 1, :)
-    x = cin .* x
+    c = 1 ./ sqrt.(degree(fg, T, dir=:in))
+    x = x .* c'
     _, x = propagate(l, fg, nothing, x, nothing, +)
-    x = cin .* x
+    x = x .* c'
     return l.σ.(l.weight * x .+ l.bias)
 end
 

test/cuda/layers/conv.jl

@@ -7,115 +7,43 @@
             0 1 0 1
             1 0 1 0]
 
-    fg = FeaturedGraph(adj, graph_type=GRAPH_T) |> gpu
-
+    fg = FeaturedGraph(adj, graph_type=GRAPH_T)
+    X = rand(Float32, in_channel, N)
+   
     @testset "GCNConv" begin
-        gc = GCNConv(in_channel=>out_channel) |> gpu
-        @test size(gc.weight) == (out_channel, in_channel)
-        @test size(gc.bias) == (out_channel,)
-        @test adjacency_matrix(gc.fg |> cpu) == adj
-
-        X = rand(in_channel, N) |> gpu
-        Y = gc(fg, X)
-        @test size(Y) == (out_channel, N)
-
-        g = Zygote.gradient(x -> sum(gc(x)), X)[1]
-        @test size(g) == size(X)
-
-        g = Zygote.gradient(model -> sum(model(X)), gc)[1]
-        @test size(g.weight) == size(gc.weight)
-        @test size(g.bias) == size(gc.bias)
+        m = GCNConv(in_channel => out_channel)
+        gpugradtest(m, fg, X)
     end
 
-
     @testset "ChebConv" begin
         k = 6
-        cc = ChebConv(in_channel=>out_channel, k) |> gpu
-        @test size(cc.weight) == (out_channel, in_channel, k)
-        @test size(cc.bias) == (out_channel,)
-        @test adjacency_matrix(cc.fg |> cpu) == adj
-        @test cc.k == k
-        
-        @test_broken begin 
-            X = rand(in_channel, N) |> gpu
-            Y = cc(fg, X)
-            @test size(Y) == (out_channel, N)
-
-            g = Zygote.gradient(x -> sum(cc(x)), X)[1]
-            @test size(g) == size(X)
-
-            g = Zygote.gradient(model -> sum(model(fg, X)), cc)[1]
-            @test size(g.weight) == size(cc.weight)
-            @test size(g.bias) == size(cc.bias)
-
-            true
-        end
+        m = ChebConv(in_channel => out_channel, k)
+        @test_broken gpugradtest(m, fg, X)
     end
 
     @testset "GraphConv" begin
-        gc = GraphConv(in_channel=>out_channel) |> gpu
-        @test size(gc.weight1) == (out_channel, in_channel)
-        @test size(gc.weight2) == (out_channel, in_channel)
-        @test size(gc.bias) == (out_channel,)
-
-        X = rand(in_channel, N) |> gpu
-        Y = gc(fg, X)
-        @test size(Y) == (out_channel, N)
-
-        g = Zygote.gradient(x -> sum(gc(fg, x)), X)[1]
-        @test size(g) == size(X)
-
-        g = Zygote.gradient(model -> sum(model(fg, X)), gc)[1]
-        @test size(g.weight1) == size(gc.weight1)
-        @test size(g.weight2) == size(gc.weight2)
-        @test size(g.bias) == size(gc.bias)
+        m = GraphConv(in_channel => out_channel)
+        gpugradtest(m, fg, X)
     end
 
     @testset "GATConv" begin
-        gat = GATConv(in_channel=>out_channel) |> gpu
-        @test size(gat.weight) == (out_channel, in_channel)
-        @test size(gat.bias) == (out_channel,)
-
-        X = rand(in_channel, N) |> gpu
-        Y = gat(fg, X)
-        @test size(Y) == (out_channel, N)
-
-        g = Zygote.gradient(x -> sum(gat(fg, x)), X)[1]
-        @test size(g) == size(X)
+        m = GATConv(in_channel => out_channel)
+        gpugradtest(m, fg, X)
+    end
 
-        g = Zygote.gradient(model -> sum(model(fg, X)), gat)[1]
-        @test size(g.weight) == size(gat.weight)
-        @test size(g.bias) == size(gat.bias)
-        @test size(g.a) == size(gat.a)
+    @testset "GINConv" begin
+        m = GINConv(Dense(in_channel, out_channel), eps=0.1f0)
+        gpugradtest(m, fg, X)
     end
 
     @testset "GatedGraphConv" begin
         num_layers = 3
-        ggc = GatedGraphConv(out_channel, num_layers) |> gpu
-        @test size(ggc.weight) == (out_channel, out_channel, num_layers)
-
-        X = rand(in_channel, N) |> gpu
-        Y = ggc(fg, X)
-        @test size(Y) == (out_channel, N)
-
-        g = Zygote.gradient(x -> sum(ggc(fg, x)), X)[1]
-        @test size(g) == size(X)
-
-        g = Zygote.gradient(model -> sum(model(fg, X)), ggc)[1]
-        @test size(g.weight) == size(ggc.weight)
+        m = GatedGraphConv(out_channel, num_layers)
+        gpugradtest(m, fg, X)
     end
 
     @testset "EdgeConv" begin
-        ec = EdgeConv(Dense(2*in_channel, out_channel)) |> gpu
-        X = rand(in_channel, N) |> gpu
-        Y = ec(fg, X)
-        @test size(Y) == (out_channel, N)
-
-        g = Zygote.gradient(x -> sum(ec(fg, x)), X)[1]
-        @test size(g) == size(X)
-
-        g = Zygote.gradient(model -> sum(model(fg, X)), ec)[1]
-        @test size(g.nn.weight) == size(ec.nn.weight)
-        @test size(g.nn.bias) == size(ec.nn.bias)
+        m = EdgeConv(Dense(2*in_channel, out_channel))
+        gpugradtest(m, fg, X)
     end
 end

test/cuda/test_utils.jl

@@ -0,0 +1,32 @@
+function gpugradtest(m, fg, x::AbstractArray{T}) where T
+    m_gpu = m |> gpu
+    fg_gpu = fg |> gpu
+    x_gpu = x |> gpu
+    ps = Flux.params(m)
+    ps_gpu = Flux.params(m_gpu)
+
+    # output
+    y = m(fg, x)
+    y_gpu = m_gpu(fg_gpu, x_gpu)
+    @test y_gpu isa CuArray{T}
+    @test Array(y_gpu) ≈ y
+
+    # input gradient
+    gs = gradient(x -> sum(m(fg, x)), x)[1]
+    gs_gpu = gradient(x_gpu -> sum(m_gpu(fg_gpu, x_gpu)), x_gpu)[1]
+    @test gs_gpu isa CuArray{T}
+    @test Array(gs_gpu) ≈ gs
+
+    # model gradient
+    gs = gradient(() -> sum(m(fg, x)), Flux.params(m))
+    gs_gpu = gradient(() -> sum(m_gpu(fg_gpu, x_gpu)), Flux.params(m_gpu))
+    for (p, p_gpu) in zip(ps, ps_gpu)
+        if gs[p] == nothing
+            @test gs_gpu[p_gpu] == nothing            
+        else
+            @test gs_gpu[p_gpu] isa CuArray{T}
+            @test Array(gs_gpu[p_gpu]) ≈ gs[p]
+        end
+    end
+    return true
+end

test/runtests.jl

@@ -11,6 +11,8 @@ using Zygote
 using Test
 CUDA.allowscalar(false)
 
+include("cuda/test_utils.jl")
+
 tests = [
     "featured_graph",
     "layers/msgpass",
@@ -22,12 +24,12 @@ tests = [
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
 # Testing all graph types. :sparse is a bit broken at the moment
-@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:dense, :coo, :sparse)
+@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo, :sparse, :dense)
     global GRAPH_T = graph_type
     for t in tests
         include("$t.jl")
 
-        if CUDA.functional() && GRAPH_T != :sparse && isdir("cuda/$t.jl")
+        if CUDA.functional() && GRAPH_T != :sparse && isfile("cuda/$t.jl")
             include("cuda/$t.jl")
         end
     end