implement new GCNConv

Author: yuehhua

Project.toml

@@ -27,13 +27,13 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 CUDA = "3"
 ChainRulesCore = "1.7"
 DataStructures = "0.18"
-FillArrays = "0.12"
+FillArrays = "0.12 - 0.13"
 Flux = "0.12"
 GraphMLDatasets = "0.1"
 GraphSignals = "0.3"
-Graphs = "1.4"
-NNlib = "0.7"
-NNlibCUDA = "0.1"
+Graphs = "1"
+NNlib = "0.7 - 0.8"
+NNlibCUDA = "0.1 - 0.2"
 Reexport = "1.1"
 Word2Vec = "0.5"
 Zygote = "0.6"

src/GeometricFlux.jl

@@ -54,8 +54,9 @@ export
     InnerProductDecoder,
     VariationalEncoder,
 
-    # layer/misc
-    Bypass,
+    # layer/utils
+    WithGraph,
+    GraphParallel,
 
     #node2vec
     node2vec
@@ -68,10 +69,10 @@ include("layers/graphlayers.jl")
 include("layers/gn.jl")
 include("layers/msgpass.jl")
 
+include("layers/utils.jl")
 include("layers/conv.jl")
 include("layers/pool.jl")
 include("models.jl")
-include("layers/misc.jl")
 
 include("sampling.jl")
 include("embedding/node2vec.jl")

src/layers/conv.jl

@@ -1,51 +1,61 @@
 """
-    GCNConv([fg,] in => out, σ=identity; bias=true, init=glorot_uniform)
+    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform)
 
-Graph convolutional layer.
+Graph convolutional layer. The input to the layer is a node feature array `X`
+of size `(num_features, num_nodes)`.
 
 # Arguments
 
-- `fg`: Optionally pass a [`FeaturedGraph`](@ref). 
 - `in`: The dimension of input features.
 - `out`: The dimension of output features.
 - `σ`: Activation function.
 - `bias`: Add learnable bias.
 - `init`: Weights' initializer.
 
+# Example
 
-The input to the layer is a node feature array `X` 
-of size `(num_features, num_nodes)`.
+```jldoctest
+julia> gc = GCNConv(1024=>256, relu)
+GCNConv(1024 => 256, relu)
+```
+
+See also [`WithGraph`](@ref) for training layer with fixed graph or subgraph.
 """
-struct GCNConv{A<:AbstractMatrix, B, F, S<:AbstractFeaturedGraph} <: AbstractGraphLayer
+struct GCNConv{A<:AbstractMatrix,B,F}
     weight::A
     bias::B
     σ::F
-    fg::S
 end
 
-function GCNConv(fg::AbstractFeaturedGraph, ch::Pair{Int,Int}, σ=identity;
+function GCNConv(ch::Pair{Int,Int}, σ=identity;
                  init=glorot_uniform, bias::Bool=true)
     in, out = ch
     W = init(out, in)
     b = Flux.create_bias(W, bias, out)
-    GCNConv(W, b, σ, fg)
+    GCNConv(W, b, σ)
 end
 
-GCNConv(ch::Pair{Int,Int}, σ = identity; kwargs...) =
-    GCNConv(NullGraph(), ch, σ; kwargs...)
-
 @functor GCNConv
 
-Flux.trainable(l::GCNConv) = (l.weight, l.bias)
+(l::GCNConv)(Ã::AbstractArray, x::AbstractArray) = l.σ.(l.weight * x * Ã .+ l.bias)
 
-function (l::GCNConv)(fg::ConcreteFeaturedGraph, x::AbstractMatrix)
+function (l::GCNConv)(fg::AbstractFeaturedGraph)
+    nf = node_feature(fg)
     Ã = Zygote.ignore() do
-        GraphSignals.normalized_adjacency_matrix(fg, eltype(x); selfloop=true)
+        GraphSignals.normalized_adjacency_matrix(fg, eltype(nf); selfloop=true)
     end
-    l.σ.(l.weight * x * Ã .+ l.bias)
+    return FeaturedGraph(fg, nf = l(Ã, nf))
 end
 
-(l::GCNConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+function (wg::WithGraph{<:GCNConv})(X::AbstractArray)
+    N = size(X, 2)
+    wg.subgraph != (:) && N != length(wg.subgraph) &&
+        throw(ArgumentError("Layer with subgraph expecting subset of features, got #V=$N but #V for subgraph $(length(wg.subgraph))."))
+    Ã = Zygote.ignore() do
+        GraphSignals.normalized_adjacency_matrix(wg.fg, eltype(X); selfloop=true)
+    end
+    return wg.layer(Ã[wg.subgraph, wg.subgraph], X)
+end
 
 function Base.show(io::IO, l::GCNConv)
     out, in = size(l.weight)

src/layers/misc.jl

@@ -0,0 +1,96 @@
+"""
+    WithGraph(layer, fg, [subgraph=:])
+
+Train GNN layers with fixed graph.
+
+# Arguments
+
+- `layer`: A GNN layer.
+- `fg`: A fixed `FeaturedGraph` to train with.
+- `subgraph`: Node indeices to get a subgraph from `fg`.
+
+# Example
+
+```jldoctest
+julia> adj = [0 1 0 1;
+              1 0 1 0;
+              0 1 0 1;
+              1 0 1 0];
+
+julia> fg = FeaturedGraph(adj);
+
+julia> gc = WithGraph(GCNConv(1024=>256), fg)
+WithGraph(GCNConv(1024 => 256), FeaturedGraph(#V=4, #E=4))
+
+julia> subgraph = [1, 2, 4]  # specify subgraph nodes
+
+julia> gc = WithGraph(GCNConv(1024=>256), fg, subgraph)
+WithGraph(GCNConv(1024 => 256), FeaturedGraph(#V=4, #E=4), subgraph=[1, 2, 4])
+```
+"""
+struct WithGraph{L,G<:AbstractFeaturedGraph,S}
+    layer::L
+    fg::G
+    subgraph::S
+end
+
+@functor WithGraph
+
+Flux.trainable(l::WithGraph) = (l.layer, )
+
+WithGraph(layer, fg::AbstractFeaturedGraph) = WithGraph(layer, fg, :)
+
+function Base.show(io::IO, l::WithGraph)
+    print(io, "WithGraph(")
+    print(io, l.layer, ", ")
+    print(io, "FeaturedGraph(#V=", nv(l.fg), ", #E=", ne(l.fg), ")")
+    l.subgraph == (:) || print(io, ", subgraph=", l.subgraph)
+    print(io, ")")
+end
+
+"""
+    GraphParallel(; node_layer=identity, edge_layer=identity, global_layer=identity)
+
+Passing features in `FeaturedGraph` in parallel. It takes `FeaturedGraph` as input
+and it can be specified by assigning layers for specific (node, edge and global) features.
+
+# Arguments
+
+- `node_layer`: A regular Flux layer for passing node features.
+- `edge_layer`: A regular Flux layer for passing edge features.
+- `global_layer`: A regular Flux layer for passing global features.
+
+# Example
+
+```jldoctest
+julia> l = GraphParallel(
+            node_layer=Dropout(0.5),
+            global_layer=Dense(10, 5)
+       )
+```
+"""
+struct GraphParallel{N,E,G}
+    node_layer::N
+    edge_layer::E
+    global_layer::G
+end
+
+@functor GraphParallel
+
+GraphParallel(; node_layer=identity, edge_layer=identity, global_layer=identity) =
+    GraphParallel(node_layer, edge_layer, global_layer)
+
+function (l::GraphParallel)(fg::FeaturedGraph)
+    nf = l.node_layer(node_feature(fg))
+    ef = l.edge_layer(edge_feature(fg))
+    gf = l.global_layer(global_feature(fg))
+    return FeaturedGraph(fg, nf=nf, ef=ef, gf=gf)
+end
+
+function Base.show(io::IO, l::GraphParallel)
+    print(io, "GraphParallel(")
+    print(io, "node_layer=", l.node_layer)
+    print(io, ", edge_layer=", l.edge_layer)
+    print(io, ", global_layer=", l.global_layer)
+    print(io, ")")
+end

src/layers/utils.jl

@@ -2,6 +2,7 @@
     T = Float32
     in_channel = 3
     out_channel = 5
+    
     N = 4
     adj = T[0 1 0 1;
            1 0 1 0;
@@ -11,17 +12,29 @@
     fg = FeaturedGraph(adj)
 
     @testset "GCNConv" begin
-        gc = GCNConv(fg, in_channel=>out_channel) |> gpu
-        @test size(gc.weight) == (out_channel, in_channel)
-        @test size(gc.bias) == (out_channel,)
-        @test collect(GraphSignals.adjacency_matrix(gc.fg)) == adj
+        X = rand(T, in_channel, N)
 
-        X = rand(in_channel, N) |> gpu
-        Y = gc(X)
-        @test size(Y) == (out_channel, N)
+        @testset "layer without graph" begin
+            gc = GCNConv(in_channel=>out_channel) |> gpu
+            @test size(gc.weight) == (out_channel, in_channel)
+            @test size(gc.bias) == (out_channel,)
 
-        g = Zygote.gradient(() -> sum(gc(X)), Flux.params(gc))
-        @test length(g.grads) == 2
+            fg = FeaturedGraph(adj, nf=X) |> gpu
+            fg_ = gc(fg)
+            @test size(node_feature(fg_)) == (out_channel, N)
+
+            g = Zygote.gradient(() -> sum(node_feature(gc(fg))), Flux.params(gc))
+            @test length(g.grads) == 4
+        end
+
+        @testset "layer with fixed graph" begin
+            gc = WithGraph(GCNConv(in_channel=>out_channel), fg) |> gpu
+            Y = gc(X |> gpu)
+            @test size(Y) == (out_channel, N)
+
+            g = Zygote.gradient(() -> sum(gc(X |> gpu)), Flux.params(gc))
+            @test length(g.grads) == 3
+        end
     end
 
 
@@ -43,19 +56,19 @@
         @test length(g.grads) == 2
     end
 
-    @testset "GraphConv" begin
-        gc = GraphConv(fg, 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,)
+    # @testset "GraphConv" begin
+    #     gc = GraphConv(fg, 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(X)
-        @test size(Y) == (out_channel, N)
+    #     X = rand(in_channel, N) |> gpu
+    #     Y = gc(X)
+    #     @test size(Y) == (out_channel, N)
 
-        g = Zygote.gradient(() -> sum(gc(X)), Flux.params(gc))
-        @test length(g.grads) == 3
-    end
+    #     g = Zygote.gradient(() -> sum(gc(X)), Flux.params(gc))
+    #     @test length(g.grads) == 3
+    # end
 
     @testset "GATConv" begin
         adj = T[1 1 0 1;

test/cuda/conv.jl

@@ -15,33 +15,16 @@
     @testset "GCNConv" begin
         X = rand(T, in_channel, N)
         Xt = transpose(rand(T, N, in_channel))
-        @testset "layer with graph" begin
-            gc = GCNConv(fg, in_channel=>out_channel)
-            @test size(gc.weight) == (out_channel, in_channel)
-            @test size(gc.bias) == (out_channel,)
-            @test GraphSignals.adjacency_matrix(gc.fg) == adj
-
-            Y = gc(X)
-            @test size(Y) == (out_channel, N)
-
-            # Test with transposed features
-            Y = gc(Xt)
-            @test size(Y) == (out_channel, N)
-
-            g = Zygote.gradient(() -> sum(gc(X)), Flux.params(gc))
-            @test length(g.grads) == 2
-        end
 
         @testset "layer without graph" begin
             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)
 
             fg = FeaturedGraph(adj, nf=X)
             fg_ = gc(fg)
             @test size(node_feature(fg_)) == (out_channel, N)
-            @test_throws ArgumentError gc(X)
+            @test_throws MethodError gc(X)
 
             # Test with transposed features
             fgt = FeaturedGraph(adj, nf=Xt)
@@ -52,6 +35,27 @@
             @test length(g.grads) == 4
         end
 
+        @testset "layer with fixed graph" begin
+            gc = WithGraph(GCNConv(in_channel=>out_channel), fg)
+            Y = gc(X)
+            @test size(Y) == (out_channel, N)
+
+            # Test with transposed features
+            Y = gc(Xt)
+            @test size(Y) == (out_channel, N)
+
+            g = Zygote.gradient(() -> sum(gc(X)), Flux.params(gc))
+            @test length(g.grads) == 3
+        end
+
+        @testset "layer with subgraph" begin
+            X = rand(T, in_channel, 3)
+            subgraph = [1,2,4]
+            gc = WithGraph(GCNConv(in_channel=>out_channel), fg, subgraph)
+            Y = gc(X)
+            @test size(Y) == (out_channel, 3)
+        end
+
         @testset "bias=false" begin
             @test length(Flux.params(GCNConv(2=>3))) == 2
             @test length(Flux.params(GCNConv(2=>3, bias=false))) == 1