Merge pull request #64 from yuehhua/develop

GCN example works again

Author: yuehhua

examples/gcn.jl

@@ -1,6 +1,7 @@
 using GeometricFlux
 using Flux
-using Flux: onehotbatch, onecold, crossentropy, throttle
+using Flux: onehotbatch, onecold, logitcrossentropy, throttle
+using Flux: @epochs
 using JLD2  # use v0.1.2
 using Statistics: mean
 using SparseArrays
@@ -31,7 +32,7 @@ model = Chain(GCNConv(adj_mat, num_features=>hidden, relu),
               softmax) |> gpu
 
 ## Loss
-loss(x, y) = crossentropy(model(x), y)
+loss(x, y) = logitcrossentropy(model(x), y)
 accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))
 
 ## Training
@@ -40,6 +41,4 @@ train_data = [(train_X, train_y)]
 opt = ADAM(0.01)
 evalcb() = @show(accuracy(train_X, train_y))
 
-for i = 1:epochs
-    Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))
-end
+@epochs epochs Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))

src/layers/conv.jl

@@ -62,7 +62,7 @@ end
 function Base.show(io::IO, l::GCNConv)
     in_channel = size(l.weight, ndims(l.weight))
     out_channel = size(l.weight, ndims(l.weight)-1)
-    print(io, "GCNConv(G(V=", nv(l.graph))
+    print(io, "GCNConv(G(V=", nv(l.fg))
     print(io, ", E), ", in_channel, "=>", out_channel)
     print(io, "GCNConv(", in_channel, "=>", out_channel)
     l.σ == identity || print(io, ", ", l.σ)

src/operations/linalg.jl

@@ -120,6 +120,11 @@ Normalized Laplacian matrix of graph `g`.
 """
 function normalized_laplacian(adj::AbstractMatrix, T::DataType=eltype(adj); selfloop::Bool=false)
     selfloop && (adj += I)
+    _normalized_laplacian(adj, T)
+end
+
+# nograd can only used without keyword arguments
+Zygote.@nograd function _normalized_laplacian(adj::AbstractMatrix, T::DataType=eltype(adj))
     inv_sqrtD = inv_sqrt_degree_matrix(adj, T, dir=:both)
     T.(I - inv_sqrtD * adj * inv_sqrtD)
 end

test/operations/linalg.jl

@@ -16,21 +16,27 @@
                    -.5 1. -.5 0.;
                    0. -.5 1. -.5;
                    -.5 0. -.5 1.]
-        scaled_lap =   [0 -0.5 0 -0.5; 
-                        -0.5 0 -0.5 -0; 
-                        0 -0.5 0 -0.5; 
+        scaled_lap =   [0 -0.5 0 -0.5;
+                        -0.5 0 -0.5 -0;
+                        0 -0.5 0 -0.5;
                         -0.5 0 -0.5 0]
 
-        for T in [Int8, Float64]
+        for T in [Int8, Int16, Int32, Int64, Float16, Float32, Float64]
             @test degree_matrix(adj, T, dir=:out) == T.(deg)
             @test degree_matrix(adj, T, dir=:out) == degree_matrix(adj, T, dir=:in)
             @test degree_matrix(adj, T, dir=:out) == degree_matrix(adj, T, dir=:both)
+            @test eltype(degree_matrix(adj, T, dir=:out)) == T
+
             @test laplacian_matrix(adj, T) == T.(lap)
+            @test eltype(laplacian_matrix(adj, T)) == T
+        end
+        for T in [Float16, Float32, Float64]
+            @test normalized_laplacian(adj, T) ≈ T.(norm_lap)
+            @test eltype(normalized_laplacian(adj, T)) == T
+            
+            @test scaled_laplacian(adj, T) ≈ T.(scaled_lap)
+            @test eltype(scaled_laplacian(adj, T)) == T
         end
-        @test normalized_laplacian(adj, Float64) ≈ norm_lap
-        @test eltype(normalized_laplacian(adj, Float32)) == Float32
-        @test scaled_laplacian(adj, Float64) ≈ scaled_lap
-        @test eltype(scaled_laplacian(adj, Float32)) == Float32
         @test neighbors(adj) == [[2,4], [1,3], [2,4], [1,3]]
     end
 
@@ -52,14 +58,21 @@
                     0 0 7 0;
                     0 0 0 4]
 
-        for T in [Int8, Float64]
+        for T in [Int8, Int16, Int32, Int64, Float16, Float32, Float64]
             @test degree_matrix(adj, T, dir=:out) == T.(deg_out)
             @test degree_matrix(adj, T, dir=:in) == T.(deg_in)
             @test degree_matrix(adj, T, dir=:both) == T.(deg_both)
+            @test eltype(degree_matrix(adj, T, dir=:out)) == T
+            @test eltype(degree_matrix(adj, T, dir=:in)) == T
+            @test eltype(degree_matrix(adj, T, dir=:both)) == T
             @test_throws DomainError degree_matrix(adj, dir=:other)
+
             @test laplacian_matrix(adj, T, dir=:out) == T.(deg_out .- adj)
             @test laplacian_matrix(adj, T, dir=:in) == T.(deg_in .- adj)
             @test laplacian_matrix(adj, T, dir=:both) == T.(deg_both .- adj)
+            @test eltype(laplacian_matrix(adj, T, dir=:out)) == T
+            @test eltype(laplacian_matrix(adj, T, dir=:in)) == T
+            @test eltype(laplacian_matrix(adj, T, dir=:both)) == T
         end
         @test neighbors(adj) == [[2,3,4], [1,3], [1,2,4], [1,3]]
     end