API design

fix

Author: yuehhua

src/GeometricFlux.jl

@@ -31,7 +31,7 @@ export
     # layers/msgpass
     MessagePassing,
 
-    # layers/conv
+    # layers/graph_conv
     GCNConv,
     ChebConv,
     GraphConv,
@@ -44,6 +44,8 @@ export
     SAGEConv,
     MeanAggregator, MeanPoolAggregator, MaxPoolAggregator,
     LSTMAggregator,
+
+    # layers/group_conv
     EEquivGraphConv,
 
     # layer/pool
@@ -73,8 +75,8 @@ include("layers/graphlayers.jl")
 include("layers/gn.jl")
 include("layers/msgpass.jl")
 
-include("layers/conv.jl")
-include("layers/groups.jl")
+include("layers/graph_conv.jl")
+include("layers/group_conv.jl")
 include("layers/pool.jl")
 include("models.jl")
 

src/layers/conv.jl renamed to src/layers/graph_conv.jl

@@ -1,45 +1,39 @@
 """
     EEquivGraphConv((in_dim, int_dim, out_dim); init)
 """
-
-struct EEquivGraphConv <: MessagePassing
-   nn_edge
-   nn_x
-   nn_h
-
-   in_dim
-   int_dim
-   out_dim
+struct EEquivGraphConv{E,X,H} <: MessagePassing
+   nn_edge::E
+   nn_x::X
+   nn_h::H
+   in_dim::Int
+   int_dim::Int
+   out_dim::Int
 end
 
 @functor EEquivGraphConv
 
-function EEquivGraphConv(dims::NTuple{3,Int}; init=glorot_uniform)
-    in_dim, int_dim, out_dim = dims
-
-    m_len = in_dim * 2 + 2
-
+function EEquivGraphConv(in_dim::Int, int_dim::Int, out_dim::Int; init=glorot_uniform)
+    m_len = 2in_dim + 2
     nn_edge = Flux.Dense(m_len, int_dim; init=init)
-
     nn_x = Flux.Dense(int_dim, 1; init=init)
     nn_h = Flux.Dense(in_dim + int_dim, out_dim; init=init)
-
-    return EEquivGraphConv(nn_edge, nn_x, nn_h, dims...)
+    return EEquivGraphConv(nn_edge, nn_x, nn_h, in_dim, int_dim, out_dim)
 end
 
-function EEquivGraphConv(nn_edge, nn_x, nn_h; init=glorot_uniform)
-
-    # Assume that these are strictly MLPs (no conv)
-    nn_edge.init(init)
-    nn_x.init(init)
-    nn_h.init(init)
-
-    in_dim = nn_edge.layers[1].W |> x->size(x)[2]
-    int_dim = nn_edge.layers[end].W |> x->size(x)[1]
-    out_dim = nn_h.layers[end].W |> x->size(x)[1]
+function EEquivGraphConv(in_dim::Int, nn_edge, nn_x, nn_h)
+    m_len = 2in_dim + 2
+    int_dim = Flux.outputsize(nn_edge, (m_len, 2))[1]
+    out_dim = Flux.outputsize(nn_h, (in_dim + int_dim, 2))[1]
     return EEquivGraphConv(nn_edge, nn_x, nn_h, in_dim, int_dim, out_dim)
 end
 
+function ϕ_edge(egnn::EEquivGraphConv, h_i, h_j, dist, a)
+    N = size(h_i, 2)
+    return egnn.nn_edge(vcat(h_i, h_j, dist, ones(N)' * a))
+end
+
+ϕ_x(egnn::EEquivGraphConv, m_ij) = egnn.nn_x(m_ij)
+
 function message(egnn::EEquivGraphConv, v_i, v_j, e)
     in_dim = egnn.in_dim
     h_i = v_i[1:in_dim,:]
@@ -56,9 +50,9 @@ function message(egnn::EEquivGraphConv, v_i, v_j, e)
         a = e[1]
     end
 
-    input = vcat(h_i, h_j, sum(abs2.(x_i - x_j); dims=1), ones(N)' * a)
-    edge_msg = egnn.nn_edge(input)
-    output_vec = vcat(edge_msg, (x_i - x_j) .* egnn.nn_x(edge_msg)[1], ones(N)')
+    dist = sum(abs2.(x_i - x_j); dims=1)
+    edge_msg = ϕ_edge(egnn, h_i, h_j, dist, a)
+    output_vec = vcat(edge_msg, (x_i - x_j) .* ϕ_x(egnn, edge_msg)[1], ones(N)')
     return reshape(output_vec, :, N)
 end
 

src/layers/groups.jl renamed to src/layers/group_conv.jl

@@ -1,4 +1,4 @@
-@testset "layer" begin
+@testset "graph conv" begin
     T = Float32
     batch_size = 10
     in_channel = 3
@@ -424,15 +424,4 @@
             end
         end
     end
-
-    @testset "EEquivGraphConv" begin
-        @testset "layer without static graph" begin
-            int_dim = 5
-            egnn = EEquivGraphConv((in_channel, int_dim, out_channel))
-
-            nf = rand(T, in_channel + 3, N)
-            fg = FeaturedGraph(adj, nf=nf)
-            fg_ = egnn(fg)
-        end
-    end
 end

test/layers/conv.jl renamed to test/layers/graph_conv.jl

@@ -0,0 +1,30 @@
+@testset "group conv" begin
+    T = Float32
+    batch_size = 10
+    in_channel = 3
+    out_channel = 5
+
+    N = 4
+    E = 4
+    adj = T[0. 1. 0. 1.;
+            1. 0. 1. 0.;
+            0. 1. 0. 1.;
+            1. 0. 1. 0.]
+    fg = FeaturedGraph(adj)
+
+    @testset "EEquivGraphConv" begin
+        @testset "layer without static graph" begin
+            int_dim = 5
+            m_len = in_channel * 2 + 2
+
+            nn_edge = Flux.Dense(m_len, int_dim)
+            nn_x = Flux.Dense(int_dim, 1)
+            nn_h = Flux.Dense(in_channel + int_dim, out_channel)
+            egnn = EEquivGraphConv(in_channel, nn_edge, nn_x, nn_h)
+
+            nf = rand(T, in_channel + 3, N)
+            fg = FeaturedGraph(adj, nf=nf)
+            fg_ = egnn(fg)
+        end
+    end
+end

test/layers/group_conv.jl

@@ -22,7 +22,8 @@ cuda_tests = [
 tests = [
      "layers/gn",
      "layers/msgpass",
-     "layers/conv",
+     "layers/graph_conv",
+     "layers/group_conv",
      "layers/pool",
      "layers/graphlayers",
      "sampling",