Merge pull request #258 from FluxML/msgpass

Fix message-passing framework

Author: yuehhua

src/GeometricFlux.jl

@@ -8,6 +8,7 @@ using Random
 using Reexport
 
 using CUDA, CUDA.CUSPARSE
+using ChainRulesCore
 using ChainRulesCore: @non_differentiable
 using FillArrays: Fill
 using Flux
@@ -19,6 +20,8 @@ using Zygote
 
 import Word2Vec: word2vec, wordvectors, get_vector
 
+const ConcreteFeaturedGraph = Union{FeaturedGraph,FeaturedSubgraph}
+
 export
     # layers/graphlayers
     AbstractGraphLayer,
@@ -50,8 +53,6 @@ export
     VGAE,
     InnerProductDecoder,
     VariationalEncoder,
-    summarize,
-    sample,
 
     # layer/selector
     bypass_graph,

src/layers/conv.jl

@@ -38,14 +38,14 @@ GCNConv(ch::Pair{Int,Int}, σ = identity; kwargs...) =
 
 Flux.trainable(l::GCNConv) = (l.weight, l.bias)
 
-function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
+function (l::GCNConv)(fg::ConcreteFeaturedGraph, x::AbstractMatrix)
     Ã = Zygote.ignore() do
         GraphSignals.normalized_adjacency_matrix(fg, eltype(x); selfloop=true)
     end
     l.σ.(l.weight * x * Ã .+ l.bias)
 end
 
-(l::GCNConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+(l::GCNConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
 
 function Base.show(io::IO, l::GCNConv)
     out, in = size(l.weight)
@@ -91,7 +91,7 @@ ChebConv(ch::Pair{Int,Int}, k::Int; kwargs...) =
 
 Flux.trainable(l::ChebConv) = (l.weight, l.bias)
 
-function (c::ChebConv)(fg::FeaturedGraph, X::AbstractMatrix{T}) where T
+function (c::ChebConv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix{T}) where T
     GraphSignals.check_num_nodes(fg, X)
     @assert size(X, 1) == size(c.weight, 2) "Input feature size must match input channel size."
 
@@ -110,7 +110,7 @@ function (c::ChebConv)(fg::FeaturedGraph, X::AbstractMatrix{T}) where T
     return Y .+ c.bias
 end
 
-(l::ChebConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+(l::ChebConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
 
 function Base.show(io::IO, l::ChebConv)
     out, in, k = size(l.weight)
@@ -165,14 +165,14 @@ message(gc::GraphConv, x_i, x_j::AbstractVector, e_ij) = gc.weight2 * x_j
 
 update(gc::GraphConv, m::AbstractVector, x::AbstractVector) = gc.σ.(gc.weight1*x .+ m .+ gc.bias)
 
-function (gc::GraphConv)(fg::FeaturedGraph, x::AbstractMatrix)
-    GraphSignals.check_num_nodes(fg, x)
-    _, x, _ = propagate(gc, graph(fg), edge_feature(fg), x, global_feature(fg), +)
+function (gc::GraphConv)(fg::ConcreteFeaturedGraph, x::AbstractMatrix)
+    # GraphSignals.check_num_nodes(fg, x)
+    _, x, _ = propagate(gc, fg, edge_feature(fg), x, global_feature(fg), +)
     x
 end
 
-(l::GraphConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
-# (l::GraphConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
+(l::GraphConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::GraphConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 function Base.show(io::IO, l::GraphConv)
     in_channel = size(l.weight1, ndims(l.weight1))
@@ -244,18 +244,22 @@ end
 
 # After some reshaping due to the multihead, we get the α from each message,
 # then get the softmax over every α, and eventually multiply the message by α
-function apply_batch_message(gat::GATConv, 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::GATConv, sg::SparseGraph, E::AbstractMatrix, X::AbstractMatrix, u)
-    @assert Zygote.ignore(() -> check_self_loops(sg)) "a vertex must have self loop (receive a message from itself)."
-    ys = map(i -> apply_batch_message(gat, i, GraphSignals.cpu_neighbors(sg, i), X), 1:nv(sg))
-    return hcat(ys...)
+function graph_attention(gat::GATConv, i, js, X::AbstractMatrix)
+    e_ij = map(j -> GeometricFlux.message(gat, _view(X, i), _view(X, j)), js)
+    E = hcat_by_sum(e_ij)
+    n = size(E, 1)
+    αs = Flux.softmax(reshape(view(E, 1, :), gat.heads, :), dims=2)
+    msgs = view(E, 2:n, :) .* reshape(αs, 1, :)
+    return reshape(msgs, (n-1)*gat.heads, :)
+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))
+    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 check_self_loops(sg::SparseGraph)
@@ -267,7 +271,7 @@ function check_self_loops(sg::SparseGraph)
     return true
 end
 
-function update_batch_vertex(gat::GATConv, M::AbstractMatrix, X::AbstractMatrix, u)
+function update_batch_vertex(gat::GATConv, ::AbstractFeaturedGraph, M::AbstractMatrix, X::AbstractMatrix, u)
     M = M .+ gat.bias
     if !gat.concat
         N = size(M, 2)
@@ -276,14 +280,14 @@ function update_batch_vertex(gat::GATConv, M::AbstractMatrix, X::AbstractMatrix,
     return M
 end
 
-function (gat::GATConv)(fg::FeaturedGraph, X::AbstractMatrix)
+function (gat::GATConv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix)
     GraphSignals.check_num_nodes(fg, X)
-    _, X, _ = propagate(gat, graph(fg), edge_feature(fg), X, global_feature(fg), +)
+    _, X, _ = propagate(gat, fg, edge_feature(fg), X, global_feature(fg), +)
     return X
 end
 
-(l::GATConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
-# (l::GATConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
+(l::GATConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::GATConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 function Base.show(io::IO, l::GATConv)
     in_channel = size(l.weight, ndims(l.weight))
@@ -335,7 +339,7 @@ message(ggc::GatedGraphConv, x_i, x_j::AbstractVector, e_ij) = x_j
 update(ggc::GatedGraphConv, m::AbstractVector, x) = m
 
 
-function (ggc::GatedGraphConv)(fg::FeaturedGraph, H::AbstractMatrix{S}) where {T<:AbstractVector,S<:Real}
+function (ggc::GatedGraphConv)(fg::ConcreteFeaturedGraph, H::AbstractMatrix{S}) where {T<:AbstractVector,S<:Real}
     GraphSignals.check_num_nodes(fg, H)
     m, n = size(H)
     @assert (m <= ggc.out_ch) "number of input features must less or equals to output features."
@@ -347,14 +351,14 @@ function (ggc::GatedGraphConv)(fg::FeaturedGraph, H::AbstractMatrix{S}) where {T
     end
     for i = 1:ggc.num_layers
         M = view(ggc.weight, :, :, i) * H
-        _, M = propagate(ggc, graph(fg), edge_feature(fg), M, global_feature(fg), +)
+        _, M = propagate(ggc, fg, edge_feature(fg), M, global_feature(fg), +)
         H, _ = ggc.gru(H, M)
     end
     H
 end
 
-(l::GatedGraphConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
-# (l::GatedGraphConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
+(l::GatedGraphConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::GatedGraphConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 
 function Base.show(io::IO, l::GatedGraphConv)
@@ -392,14 +396,14 @@ Flux.trainable(l::EdgeConv) = (l.nn,)
 message(ec::EdgeConv, x_i::AbstractVector, x_j::AbstractVector, e_ij) = ec.nn(vcat(x_i, x_j .- x_i))
 update(ec::EdgeConv, m::AbstractVector, x) = m
 
-function (ec::EdgeConv)(fg::FeaturedGraph, X::AbstractMatrix)
+function (ec::EdgeConv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix)
     GraphSignals.check_num_nodes(fg, X)
-    _, X, _ = propagate(ec, graph(fg), edge_feature(fg), X, global_feature(fg), ec.aggr)
+    _, X, _ = propagate(ec, fg, edge_feature(fg), X, global_feature(fg), ec.aggr)
     X
 end
 
-(l::EdgeConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
-# (l::EdgeConv)(fg::FeaturedGraph) = propagate(l, fg, l.aggr)  # edge number check break this
+(l::EdgeConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::EdgeConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, l.aggr)  # edge number check break this
 
 function Base.show(io::IO, l::EdgeConv)
     print(io, "EdgeConv(", l.nn)
@@ -443,15 +447,15 @@ Flux.trainable(g::GINConv) = (fg=g.fg, nn=g.nn)
 message(g::GINConv, x_i::AbstractVector, x_j::AbstractVector) = x_j 
 update(g::GINConv, m::AbstractVector, x) = g.nn((1 + g.eps) * x + m)
 
-function (g::GINConv)(fg::FeaturedGraph, X::AbstractMatrix)
+function (g::GINConv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix)
     gf = graph(fg)
     GraphSignals.check_num_nodes(gf, X)
-    _, X, _ = propagate(g, graph(fg), edge_feature(fg), X, global_feature(fg), +)
+    _, X, _ = propagate(g, fg, edge_feature(fg), X, global_feature(fg), +)
     X
 end
 
-(l::GINConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
-# (l::GINConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
+(l::GINConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::GINConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 
 """
@@ -512,17 +516,17 @@ message(c::CGConv,
 end
 update(c::CGConv, m::AbstractVector, x) = x + m
 
-function (c::CGConv)(fg::FeaturedGraph, X::AbstractMatrix, E::AbstractMatrix)
+function (c::CGConv)(fg::ConcreteFeaturedGraph, X::AbstractMatrix, E::AbstractMatrix)
     GraphSignals.check_num_nodes(fg, X)
     GraphSignals.check_num_edges(fg, E)
-    _, Y, _ = propagate(c, graph(fg), E, X, global_feature(fg), +)
+    _, Y, _ = propagate(c, fg, E, X, global_feature(fg), +)
     Y
 end
 
-(l::CGConv)(fg::FeaturedGraph) = FeaturedGraph(fg,
+(l::CGConv)(fg::AbstractFeaturedGraph) = FeaturedGraph(fg,
                                                nf=l(fg, node_feature(fg), edge_feature(fg)),
                                                ef=edge_feature(fg))
-# (l::CGConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
+# (l::CGConv)(fg::AbstractFeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 (l::CGConv)(X::AbstractMatrix, E::AbstractMatrix) = l(l.fg, X, E)
 

src/layers/gn.jl

@@ -1,7 +1,23 @@
-_view(::Nothing, i) = nothing
-_view(A::Fill{T,2,Axes}, i) where {T,Axes} = view(A, :, 1)
+_view(::Nothing, idx) = nothing
+_view(A::Fill{T,2,Axes}, idx) where {T,Axes} = fill(A.value, A.axes[1], length(idx))
+
+function _view(A::SubArray{T,2,S}, idx) where {T,S<:Fill}
+    p = parent(A)
+    return Fill(p.value, p.axes[1].stop, length(idx))
+end
+
 _view(A::AbstractMatrix, idx) = view(A, :, idx)
 
+function _view(A::SubArray{T,2,S}, idxs) where {T,S<:AbstractMatrix}
+    view_idx = A.indices[2]
+    if view_idx == idxs
+        return A
+    else
+        idxs = findall(x -> x in idxs, view_idx)
+        return view(A, :, idxs)
+    end
+end
+
 aggregate(aggr::typeof(+), X) = vec(sum(X, dims=2))
 aggregate(aggr::typeof(-), X) = -vec(sum(X, dims=2))
 aggregate(aggr::typeof(*), X) = vec(prod(X, dims=2))
@@ -16,43 +32,48 @@ abstract type GraphNet <: AbstractGraphLayer end
 @inline update_vertex(gn::GraphNet, ē, vi, u) = vi
 @inline update_global(gn::GraphNet, ē, v̄, u) = u
 
-@inline function update_batch_edge(gn::GraphNet, sg::SparseGraph, E, V, u)
-    ys = map(i -> apply_batch_message(gn, sg, i, GraphSignals.cpu_neighbors(sg, i), E, V, u), vertices(sg))
-    return hcat(ys...)
+@inline function update_batch_edge(gn::GraphNet, fg::AbstractFeaturedGraph, E, V, u)
+    es = Zygote.ignore(()->cpu(GraphSignals.incident_edges(fg)))
+    xs = Zygote.ignore(()->cpu(GraphSignals.repeat_nodes(fg)))
+    nbrs = Zygote.ignore(()->cpu(GraphSignals.neighbors(fg)))
+    ms = map((e,i,j)->update_edge(gn, _view(E, e), _view(V, i), _view(V, j), u), es, xs, nbrs)
+    M = hcat_by_sum(ms)
+    return M
 end
 
-@inline function apply_batch_message(gn::GraphNet, sg::SparseGraph, i, js, E, V, u)
-    # js still CuArray
-    es = Zygote.ignore(() -> GraphSignals.cpu_incident_edges(sg, i))
-    ys = map(k -> update_edge(gn, _view(E, es[k]), _view(V, i), _view(V, js[k]), u), 1:length(js))
-    return hcat(ys...)
+@inline function update_batch_vertex(gn::GraphNet, fg::AbstractFeaturedGraph, Ē, V, u)
+    nodes = Zygote.ignore(()->vertices(fg))
+    vs = map(n->update_vertex(gn, _view(Ē, n), _view(V, n), u), nodes)
+    V_ = hcat_by_sum(vs)
+    return V_
 end
 
-@inline function update_batch_vertex(gn::GraphNet, Ē, V, u)
-    ys = map(i -> update_vertex(gn, _view(Ē, i), _view(V, i), u), 1:size(V,2))
-    return hcat(ys...)
+@inline function aggregate_neighbors(gn::GraphNet, fg::AbstractFeaturedGraph, aggr, E)
+    N = nv(parent(fg))
+    xs = Zygote.ignore(()->cpu(GraphSignals.repeat_nodes(fg)))
+    Ē = NNlib.scatter(aggr, E, xs; dstsize=(size(E, 1), N))
+    return Ē
 end
-
-@inline aggregate_neighbors(gn::GraphNet, sg::SparseGraph, aggr, E) = neighbor_scatter(aggr, E, sg)
-@inline aggregate_neighbors(gn::GraphNet, sg::SparseGraph, aggr::Nothing, @nospecialize E) = nothing
+@inline aggregate_neighbors(gn::GraphNet, fg::AbstractFeaturedGraph, aggr::Nothing, @nospecialize E) = nothing
 
 @inline aggregate_edges(gn::GraphNet, aggr, E) = aggregate(aggr, E)
 @inline aggregate_edges(gn::GraphNet, aggr::Nothing, @nospecialize E) = nothing
 
 @inline aggregate_vertices(gn::GraphNet, aggr, V) = aggregate(aggr, V)
 @inline aggregate_vertices(gn::GraphNet, aggr::Nothing, @nospecialize V) = nothing
 
-function propagate(gn::GraphNet, fg::FeaturedGraph, naggr=nothing, eaggr=nothing, vaggr=nothing)
-    E, V, u = propagate(gn, graph(fg), edge_feature(fg), node_feature(fg), global_feature(fg), naggr, eaggr, vaggr)
+function propagate(gn::GraphNet, fg::AbstractFeaturedGraph, naggr=nothing, eaggr=nothing, vaggr=nothing)
+    E, V, u = propagate(gn, fg, edge_feature(fg), node_feature(fg), global_feature(fg), naggr, eaggr, vaggr)
     FeaturedGraph(fg, nf=V, ef=E, gf=u)
 end
 
-function propagate(gn::GraphNet, sg::SparseGraph, E, V, u, naggr=nothing, eaggr=nothing, vaggr=nothing)
-    E = update_batch_edge(gn, sg, E, V, u)
-    Ē = aggregate_neighbors(gn, sg, naggr, E)
-    V = update_batch_vertex(gn, Ē, V, u)
+function propagate(gn::GraphNet, fg::AbstractFeaturedGraph, E::AbstractArray, V::AbstractArray, u::AbstractArray,
+                   naggr=nothing, eaggr=nothing, vaggr=nothing)
+    E = update_batch_edge(gn, fg, E, V, u)
+    Ē = aggregate_neighbors(gn, fg, naggr, E)
+    V = update_batch_vertex(gn, fg, Ē, V, u)
     ē = aggregate_edges(gn, eaggr, E)
     v̄ = aggregate_vertices(gn, vaggr, V)
     u = update_global(gn, ē, v̄, u)
-    return E, V, u
+    return parent(E), parent(V), u
 end

src/layers/graphlayers.jl

@@ -1,3 +1,4 @@
 abstract type AbstractGraphLayer end
 
 (l::AbstractGraphLayer)(x::AbstractMatrix) = l(l.fg, x)
+(l::AbstractGraphLayer)(::NullGraph, x::AbstractMatrix) = throw(ArgumentError("concrete FeaturedGraph is not provided."))

src/utils.jl

@@ -1,22 +1,31 @@
-"""
-    accumulated_edges(adj)
+function hcat_by_sum(xs::AbstractVector)
+    T = eltype(xs[1])
+    dim = size(xs[1], 1)
+    N = length(xs)
 
-Return a vector which acts as a mapping table. The index is the vertex index,
-value is accumulated numbers of edge (current vertex not included).
-"""
-accumulated_edges(adj::AbstractVector{<:AbstractVector{<:Integer}}) = [0, cumsum(map(length, adj))...]
+    ns = map(x->size(x,2), xs)
+    pushfirst!(ns, 1)
+    cumsum!(ns, ns)
 
-function generate_cluster(M::AbstractArray{T,N}, accu_edge) where {T,N}
-    num_V = length(accu_edge) - 1
-    num_E = accu_edge[end]
-    cluster = similar(M, Int, num_E)
-    @inbounds for i = 1:num_V
-        j = accu_edge[i]
-        k = accu_edge[i+1]
-        cluster[j+1:k] .= i
+    A = similar(xs[1], T, dim, ns[end]-1)
+    for i in 1:N
+        A[:, ns[i]:(ns[i+1]-1)] .= xs[i]
     end
-    cluster
+    return A
 end
 
-@non_differentiable accumulated_edges(x...)
-@non_differentiable generate_cluster(x...)
+function ChainRulesCore.rrule(::typeof(hcat_by_sum), xs::AbstractVector)
+    N = length(xs)
+
+    ns = map(x->size(x,2), xs)
+    pushfirst!(ns, 1)
+    cumsum!(ns, ns)
+
+    hcat_by_sum_pullback(Δ) = (NoTangent(), ntuple(i->view(Δ,:,ns[i]:(ns[i+1]-1)), N))
+    hcat_by_sum(xs), hcat_by_sum_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(parent), A::Base.SubArray)
+    parent_pullback(Δ) = (NoTangent(), view(Δ, A.indices...))
+    parent(A), parent_pullback
+end