Merge pull request #259 from abieler/gatv2

Adds GATv2 layer

Author: yuehhua

src/GeometricFlux.jl

@@ -35,6 +35,7 @@ export
     ChebConv,
     GraphConv,
     GATConv,
+    GATv2Conv,
     GatedGraphConv,
     EdgeConv,
     GINConv,

src/layers/conv.jl

@@ -272,7 +272,7 @@ end
 
 function update_batch_edge(gat::GATConv, fg::AbstractFeaturedGraph, E::AbstractMatrix, X::AbstractMatrix, u)
     @assert Zygote.ignore(() -> check_self_loops(graph(fg))) "a vertex must have self loop (receive a message from itself)."
-    nodes = Zygote.ignore(()->vertices(fg))
+    nodes = Zygote.ignore(()->vertices(graph(fg)))
     nbr = i->cpu(GraphSignals.neighbors(graph(fg), i))
     ms = map(i -> graph_attention(gat, i, Zygote.ignore(()->nbr(i)), X), nodes)
     M = hcat_by_sum(ms)
@@ -292,7 +292,7 @@ function update_batch_vertex(gat::GATConv, ::AbstractFeaturedGraph, M::AbstractM
     M = M .+ gat.bias
     if !gat.concat
         N = size(M, 2)
-        M = reshape(mean(reshape(M, :, gat.heads, N), dims=2), :, N)
+        M = reshape(mean(reshape(M, gat.heads, :, N), dims=1), :, N)
     end
     return M
 end
@@ -315,6 +315,102 @@ function Base.show(io::IO, l::GATConv)
 end
 
 
+"""
+    GATv2Conv([fg,] in => out;
+            heads=1,
+            concat=true,
+            init=glorot_uniform    
+            negative_slope=0.2)
+
+GATv2 Layer as introduced in https://arxiv.org/abs/2105.14491
+
+# Arguments
+
+- `fg`: Optionally pass a [`FeaturedGraph`](@ref). 
+- `in`: The dimension of input features.
+- `out`: The dimension of output features.
+- `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.
+"""
+struct GATv2Conv{V<:AbstractFeaturedGraph, T, A<:AbstractMatrix{T}, B} <: MessagePassing
+    fg::V
+    wi::A
+    wj::A
+    biasi::B
+    biasj::B
+    a::A
+    negative_slope::T
+    channel::Pair{Int, Int}
+    heads::Int
+    concat::Bool
+end
+
+function GATv2Conv(
+    fg::AbstractFeaturedGraph,
+    ch::Pair{Int,Int};
+    heads::Int=1,
+    concat::Bool=true,
+    negative_slope=0.2f0,
+    bias::Bool=true,
+    init=glorot_uniform,
+)
+    in, out = ch
+    wi = init(out*heads, in)
+    wj = init(out*heads, in)
+    bi = Flux.create_bias(wi, bias, out*heads)
+    bj = Flux.create_bias(wj, bias, out*heads)
+    a = init(out, heads)
+    GATv2Conv(fg, wi, wj, bi, bj, a, negative_slope, ch, heads, concat)
+end
+
+GATv2Conv(ch::Pair{Int,Int}; kwargs...) = GATv2Conv(NullGraph(), ch; kwargs...)
+
+@functor GATv2Conv
+
+Flux.trainable(l::GATv2Conv) = (l.wi, l.wj, l.biasi, l.biasj, l.a)
+
+function message(gat::GATv2Conv, x_i::AbstractVector, x_j::AbstractVector)
+    xi = reshape(gat.wi * x_i + gat.biasi, :, gat.heads)
+    xj = reshape(gat.wj * x_j + gat.biasj, :, gat.heads)
+    eij = gat.a' * leakyrelu.(xi + xj, gat.negative_slope)
+    vcat(eij, xj)
+end
+
+function graph_attention(gat::GATv2Conv, i, js, X::AbstractMatrix)
+    e_ij = mapreduce(j -> GeometricFlux.message(gat, _view(X, i), _view(X, j)), hcat, js)
+    n = size(e_ij, 1)
+    αs = Flux.softmax(reshape(view(e_ij, 1, :), gat.heads, :), dims=2)
+    msgs = view(e_ij, 2:n, :) .* reshape(αs, 1, :)
+    reshape(msgs, (n-1)*gat.heads, :)
+end
+
+function update_batch_edge(gat::GATv2Conv, fg::AbstractFeaturedGraph, E::AbstractMatrix, X::AbstractMatrix, u)
+    @assert Zygote.ignore(() -> check_self_loops(graph(fg))) "a vertex must have self loop (receive a message from itself)."
+    nodes = Zygote.ignore(()->vertices(graph(fg)))
+    nbr = i->cpu(GraphSignals.neighbors(graph(fg), i))
+    ms = map(i -> graph_attention(gat, i, Zygote.ignore(()->nbr(i)), X), nodes)
+    M = hcat_by_sum(ms)
+    return M
+end
+
+function update_batch_vertex(gat::GATv2Conv, ::AbstractFeaturedGraph, M::AbstractMatrix, X::AbstractMatrix, u)
+    if !gat.concat
+        N = size(M, 2)
+        M = reshape(mean(reshape(M, gat.heads, :, N), dims=1), :, N)
+    end
+    return M
+end
+
+function (gat::GATv2Conv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix)
+    GraphSignals.check_num_nodes(fg, X)
+    _, X, _ = propagate(gat, fg, edge_feature(fg), X, global_feature(fg), +)
+    return X
+end
+
+(l::GATv2Conv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+
+
 """
     GatedGraphConv([fg,] out, num_layers; aggr=+, init=glorot_uniform)
 

test/layers/conv.jl

@@ -227,6 +227,77 @@
         end
     end
 
+    @testset "GATv2Conv" begin
+        adj1 = [1 1 0 1;
+                1 1 1 0;
+                0 1 1 1;
+                1 0 1 1]
+        fg1 = FeaturedGraph(adj1)
+
+        # isolated_vertex
+        adj2 = [1 0 0 1;
+                0 1 0 0;
+                0 0 1 1;
+                1 0 1 1]
+        fg2 = FeaturedGraph(adj2)
+
+        X = rand(T, in_channel, N)
+        Xt = transpose(rand(T, N, in_channel))
+
+        @testset "layer with graph" begin
+            for heads = [1, 2], concat = [true, false], adj_gat in [adj1, adj2]
+                fg_gat = FeaturedGraph(adj_gat)
+                gat2 = GATv2Conv(fg_gat, in_channel=>out_channel, heads=heads, concat=concat)
+
+                @test size(gat2.wi) == (out_channel * heads, in_channel)
+                @test size(gat2.wi) == (out_channel * heads, in_channel)
+                @test size(gat2.biasi) == (out_channel * heads,)
+                @test size(gat2.biasj) == (out_channel * heads,)
+                @test size(gat2.a) == (out_channel, heads)
+
+                Y = gat2(X)
+                @test size(Y) == (concat ? (out_channel*heads, N) : (out_channel, N))
+
+                # Test with transposed features
+                Y = gat2(Xt)
+                @test size(Y) == (concat ? (out_channel*heads, N) : (out_channel, N))
+
+                g = Zygote.gradient(() -> sum(gat2(X)), Flux.params(gat2))
+                @test length(g.grads) == 5
+            end
+        end
+
+        @testset "layer without graph" begin
+            for heads = [1, 2], concat = [true, false], adj_gat in [adj1, adj2]
+                fg_gat = FeaturedGraph(adj_gat, nf=X)
+                gat2 = GATv2Conv(in_channel=>out_channel, heads=heads, concat=concat)
+                @test size(gat2.wi) == (out_channel * heads, in_channel)
+                @test size(gat2.wi) == (out_channel * heads, in_channel)
+                @test size(gat2.biasi) == (out_channel * heads,)
+                @test size(gat2.biasj) == (out_channel * heads,)
+                @test size(gat2.a) == (out_channel, heads)
+
+                fg_ = gat2(fg_gat)
+                Y = node_feature(fg_)
+                @test size(Y) == (concat ? (out_channel*heads, N) : (out_channel, N))
+                @test_throws ArgumentError gat2(X)
+
+                # Test with transposed features
+                fgt = FeaturedGraph(adj_gat, nf=Xt)
+                fgt_ = gat2(fgt)
+                @test size(node_feature(fgt_)) == (concat ? (out_channel*heads, N) : (out_channel, N))
+
+                g = Zygote.gradient(() -> sum(node_feature(gat2(fg_gat))), Flux.params(gat2))
+                @test length(g.grads) == 7
+            end
+        end
+
+        @testset "bias=false" begin
+            @test length(Flux.params(GATv2Conv(2=>3))) == 5
+            @test length(Flux.params(GATv2Conv(2=>3, bias=false))) == 3
+        end
+    end
+
     @testset "GatedGraphConv" begin
         num_layers = 3
         X = rand(T, in_channel, N)