Merge pull request #340 from FluxML/refactor

Fix tests and refactor

Author: yuehhua

src/layers/graph_conv.jl

@@ -128,7 +128,7 @@ function (l::ChebConv)(fg::AbstractFeaturedGraph)
     nf = node_feature(fg)
     GraphSignals.check_num_nodes(fg, nf)
     @assert size(nf, 1) == size(l.weight, 2) "Input feature size must match input channel size."
-    
+
     L̃ = ChainRulesCore.ignore_derivatives() do
         GraphSignals.scaled_laplacian(fg, eltype(nf))
     end
@@ -245,7 +245,7 @@ Graph attentional layer.
 - `out`: The dimension of output features.
 - `bias::Bool`: Keyword argument, whether to learn the additive bias.
 - `σ`: Activation function.
-- `heads`: Number attention heads 
+- `heads`: Number attention heads
 - `concat`: Concatenate layer output or not. If not, layer output is averaged.
 - `negative_slope::Real`: Keyword argument, the parameter of LeakyReLU.
 
@@ -280,7 +280,7 @@ end
 
 function GATConv(ch::Pair{Int,Int}, σ=identity; heads::Int=1, concat::Bool=true,
                  negative_slope=0.2f0, init=glorot_uniform, bias::Bool=true)
-    in, out = ch             
+    in, out = ch
     W = init(out*heads, in)
     b = Flux.create_bias(W, bias, out*heads)
     a = init(2*out, heads)
@@ -372,7 +372,7 @@ Graph attentional layer v2.
 - `in`: The dimension of input features.
 - `out`: The dimension of output features.
 - `σ`: Activation function.
-- `heads`: Number attention heads 
+- `heads`: Number attention heads
 - `concat`: Concatenate layer output or not. If not, layer output is averaged.
 - `negative_slope::Real`: Keyword argument, the parameter of LeakyReLU.
 
@@ -661,7 +661,7 @@ GINConv(nn, eps=0f0) = GINConv(nn, eps)
 
 Flux.trainable(g::GINConv) = (g.nn,)
 
-message(g::GINConv, x_i::AbstractArray, x_j::AbstractArray) = x_j 
+message(g::GINConv, x_i::AbstractArray, x_j::AbstractArray) = x_j
 update(g::GINConv, m::AbstractArray, x::AbstractArray) = g.nn((1 + g.eps) * x + m)
 
 # For variable graph
@@ -705,30 +705,24 @@ CGConv(node dim=128, edge dim=32)
 
 See also [`WithGraph`](@ref) for training layer with static graph.
 """
-struct CGConv{A<:AbstractMatrix,B} <: MessagePassing
-    Wf::A
-    Ws::A
-    bf::B
-    bs::B
+struct CGConv{A,B} <: MessagePassing
+    f::A
+    s::B
 end
 
 @functor CGConv
 
-Flux.trainable(l::CGConv) = (l.Wf, l.Ws, l.bf, l.bs)
-
 function CGConv(dims::NTuple{2,Int}; init=glorot_uniform, bias=true)
     node_dim, edge_dim = dims
-    Wf = init(node_dim, 2*node_dim + edge_dim)
-    Ws = init(node_dim, 2*node_dim + edge_dim)
-    bf = Flux.create_bias(Wf, bias, node_dim)
-    bs = Flux.create_bias(Ws, bias, node_dim)
-    return CGConv(Wf, Ws, bf, bs)
+    f = Dense(2*node_dim + edge_dim, node_dim; bias=bias, init=init)
+    s = Dense(2*node_dim + edge_dim, node_dim; bias=bias, init=init)
+    return CGConv(f, s)
 end
 
-function message(c::CGConv, x_i::AbstractArray, x_j::AbstractArray, e::AbstractArray)
+function message(l::CGConv, x_i::AbstractArray, x_j::AbstractArray, e::AbstractArray)
     z = vcat(x_i, x_j, e)
 
-    return σ.(_matmul(c.Wf, z) .+ c.bf) .* softplus.(_matmul(c.Ws, z) .+ c.bs)
+    return σ.(l.f(z)) .* softplus.(l.s(z))
 end
 
 update(c::CGConv, m::AbstractArray, x) = x + m
@@ -752,7 +746,7 @@ function (l::CGConv)(el::NamedTuple, X::AbstractArray, E::AbstractArray)
 end
 
 function Base.show(io::IO, l::CGConv)
-    node_dim, d = size(l.Wf)
+    node_dim, d = size(l.f.weight)
     edge_dim = d - 2*node_dim
     print(io, "CGConv(node dim=", node_dim, ", edge dim=", edge_dim, ")")
 end

test/layers/graph_conv.jl

@@ -74,7 +74,7 @@
             @test size(cc.weight) == (out_channel, in_channel, k)
             @test size(cc.bias) == (out_channel,)
             @test cc.k == k
-            
+
             fg = FeaturedGraph(adj, nf=X)
             fg_ = cc(fg)
             @test size(node_feature(fg_)) == (out_channel, N)
@@ -90,7 +90,7 @@
         end
 
         @testset "layer with static graph" begin
-            cc = WithGraph(fg, ChebConv(in_channel=>out_channel, k))            
+            cc = WithGraph(fg, ChebConv(in_channel=>out_channel, k))
             Y = cc(X)
             @test size(Y) == (out_channel, N)
 
@@ -286,7 +286,7 @@
             g = gradient(() -> sum(node_feature(ec(fg))), Flux.params(ec))
             @test length(g.grads) == 4
         end
-        
+
         @testset "layer with static graph" begin
             X = rand(T, in_channel, N, batch_size)
             ec = WithGraph(fg, EdgeConv(Dense(2*in_channel, out_channel)))
@@ -322,10 +322,8 @@
         nn = Flux.Chain(Dense(in_channel, out_channel))
         eps = 0.001
         @testset "layer without graph" begin
-            gc = GraphConv(in_channel=>out_channel)
-            @test size(gc.weight1) == (out_channel, in_channel)
-            @test size(gc.weight2) == (out_channel, in_channel)
-            @test size(gc.bias) == (out_channel,)
+            gc = GINConv(nn, eps)
+            @test gc.nn == nn
 
             X = rand(T, in_channel, N)
             fg = FeaturedGraph(adj, nf=X)
@@ -334,7 +332,7 @@
             @test_throws MethodError gc(X)
 
             g = gradient(() -> sum(node_feature(gc(fg))), Flux.params(gc))
-            @test length(g.grads) == 5
+            @test length(g.grads) == 4
         end
 
         @testset "layer with static graph" begin
@@ -351,10 +349,10 @@
     @testset "CGConv" begin
         @testset "layer without graph" begin
             cgc = CGConv((in_channel, in_channel_edge))
-            @test size(cgc.Wf) == (in_channel, 2 * in_channel + in_channel_edge)
-            @test size(cgc.Ws) == (in_channel, 2 * in_channel + in_channel_edge)
-            @test size(cgc.bf) == (in_channel,)
-            @test size(cgc.bs) == (in_channel,)
+            @test size(cgc.f.weight) == (in_channel, 2 * in_channel + in_channel_edge)
+            @test size(cgc.s.weight) == (in_channel, 2 * in_channel + in_channel_edge)
+            @test size(cgc.f.bias) == (in_channel,)
+            @test size(cgc.s.bias) == (in_channel,)
 
             nf = rand(T, in_channel, N)
             ef = rand(T, in_channel_edge, E)
@@ -403,7 +401,7 @@
                 end
             end
         end
-        
+
         @testset "layer with static graph" begin
             for conv in aggregators
                 X = rand(T, in_channel, N, batch_size)