Merge pull request #108 from yuehhua/develop

Refactor

Author: yuehhua

src/layers/conv.jl

@@ -27,16 +27,16 @@ struct GCNConv{T,F,S<:AbstractFeaturedGraph}
 end
 
 function GCNConv(ch::Pair{<:Integer,<:Integer}, σ = identity;
-                 init=glorot_uniform, T::DataType=Float32, bias::Bool=true, cache::Bool=true)
+                 init=glorot_uniform, T::DataType=Float32, bias::Bool=true)
     b = bias ? T.(init(ch[2])) : zeros(T, ch[2])
-    fg = cache ? FeaturedGraph() : NullGraph()
+    fg = NullGraph()
     GCNConv(T.(init(ch[2], ch[1])), b, σ, fg)
 end
 
 function GCNConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, σ = identity;
-                 init=glorot_uniform, T::DataType=Float32, bias::Bool=true, cache::Bool=true)
+                 init=glorot_uniform, T::DataType=Float32, bias::Bool=true)
     b = bias ? T.(init(ch[2])) : zeros(T, ch[2])
-    fg = cache ? FeaturedGraph(adj) : NullGraph()
+    fg = FeaturedGraph(adj)
     GCNConv(T.(init(ch[2], ch[1])), b, σ, fg)
 end
 
@@ -56,8 +56,10 @@ end
 
 function (g::GCNConv)(fg::FeaturedGraph)
     X = node_feature(fg)
-    A = adjacency_matrix(fg)
-    g.fg isa NullGraph || (g.fg.graph = A)
+    A = adjacency_matrix(fg) # TODO: choose graph from g or fg
+    Zygote.ignore() do
+        g.fg isa NullGraph || (g.fg.graph = A)
+    end
     X_ = g(A, X)
     FeaturedGraph(A, X_)
 end
@@ -97,16 +99,16 @@ struct ChebConv{T,S<:AbstractFeaturedGraph}
 end
 
 function ChebConv(adj::AbstractMatrix, ch::Pair{<:Integer,<:Integer}, k::Integer;
-                  init = glorot_uniform, T::DataType=Float32, bias::Bool=true, cache::Bool=true)
+                  init = glorot_uniform, T::DataType=Float32, bias::Bool=true)
     b = bias ? init(ch[2]) : zeros(T, ch[2])
-    fg = cache ? FeaturedGraph(adj) : NullGraph()
+    fg = FeaturedGraph(adj)
     ChebConv(init(ch[2], ch[1], k), b, fg, k, ch[1], ch[2])
 end
 
 function ChebConv(ch::Pair{<:Integer,<:Integer}, k::Integer;
-                  init = glorot_uniform, T::DataType=Float32, bias::Bool=true, cache::Bool=true)
+                  init = glorot_uniform, T::DataType=Float32, bias::Bool=true)
     b = bias ? init(ch[2]) : zeros(T, ch[2])
-    fg = cache ? FeaturedGraph() : NullGraph()
+    fg = NullGraph()
     ChebConv(init(ch[2], ch[1], k), b, fg, k, ch[1], ch[2])
 end
 
@@ -138,7 +140,9 @@ end
 function (c::ChebConv)(fg::FeaturedGraph)
     @assert has_graph(fg) "A given FeaturedGraph must contain a graph."
     g = graph(fg)
-    c.fg isa NullGraph || (c.fg.graph = g)
+    Zygote.ignore() do
+        c.fg isa NullGraph || (c.fg.graph = g)
+    end
     X = node_feature(fg)
     L̃ = scaled_laplacian(adjacency_matrix(fg))
     L̃ = convert(typeof(X), L̃)

src/layers/gn.jl

@@ -8,24 +8,19 @@ abstract type GraphNet end
 @inline update_global(gn::T, ē, v̄, u) where {T<:GraphNet} = u
 
 @inline function update_batch_edge(gn::T, adj, E, V, u) where {T<:GraphNet}
+    n = size(adj, 1)
     edge_idx = edge_index_table(adj)
-    E_ = Vector[]
-    for (i, js) = enumerate(adj)
-        for j = js
-            k = edge_idx[(i,j)]
-            e = update_edge(gn, get_feature(E, k), get_feature(V, i), get_feature(V, j), u)
-            push!(E_, e)
-        end
-    end
+    E_ = [_apply_batch_message(gn, i, adj[i], edge_idx, E, V, u) for i in 1:n]
+    hcat(E_...)
+end
+
+@inline function _apply_batch_message(gn::T, i, js, edge_idx, E, V, u) where {T<:GraphNet}
+    E_ = [update_edge(gn, get_feature(E, edge_idx[(i,j)]), get_feature(V, i), get_feature(V, j), u) for j = js]
     hcat(E_...)
 end
 
 @inline function update_batch_vertex(gn::T, Ē, V, u) where {T<:GraphNet}
-    V_ = Vector[]
-    for i = 1:size(V,2)
-        v = update_vertex(gn, get_feature(Ē, i), get_feature(V, i), u)
-        push!(V_, v)
-    end
+    V_ = [update_vertex(gn, get_feature(Ē, i), get_feature(V, i), u) for i = 1:size(V,2)]
     hcat(V_...)
 end
 

src/models.jl

@@ -96,8 +96,8 @@ end
 
 function VariationalEncoder(nn, h_dim::Integer, z_dim::Integer)
     VariationalEncoder(nn,
-                       GCNConv(h_dim=>z_dim; cache=false),
-                       GCNConv(h_dim=>z_dim; cache=false),
+                       GCNConv(h_dim=>z_dim),
+                       GCNConv(h_dim=>z_dim),
                        z_dim)
 end
 

test/layers/conv.jl

@@ -26,7 +26,7 @@ adj = [0. 1. 0. 1.;
         end
 
         @testset "layer without graph" begin
-            gc = GCNConv(in_channel=>out_channel, cache=false)
+            gc = GCNConv(in_channel=>out_channel)
             @test size(gc.weight) == (out_channel, in_channel)
             @test size(gc.bias) == (out_channel,)
             @test !has_graph(gc.fg)
@@ -62,7 +62,7 @@ adj = [0. 1. 0. 1.;
         end
 
         @testset "layer without graph" begin
-            cc = ChebConv(in_channel=>out_channel, k, cache=false)
+            cc = ChebConv(in_channel=>out_channel, k)
             @test size(cc.weight) == (out_channel, in_channel, k)
             @test size(cc.bias) == (out_channel,)
             @test !has_graph(cc.fg)