simplify MessagePassing methods

fix bug for GATConv and ChebConv

Author: yuehhua

Project.toml

@@ -5,6 +5,7 @@ version = "0.7.7"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
@@ -20,9 +21,10 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-CUDA = "3.3"
+CUDA = "3"
+ChainRulesCore = "1.7"
 DataStructures = "0.18"
-FillArrays = "0.11, 0.12"
+FillArrays = "0.12"
 Flux = "0.12"
 GraphMLDatasets = "0.1"
 GraphSignals = "0.3"

docs/make.jl

@@ -13,7 +13,8 @@ makedocs(
     pages = ["Home" => "index.md",
              "Get started" => "start.md",
              "Basics" =>
-               ["Building layers" => "basics/layers.md",
+               ["Graph convolutions" => "basics/conv.md",
+                "Building layers" => "basics/layers.md",
                 "Graph passing" => "basics/passgraph.md"],
              "Cooperate with Flux layers" => "cooperate.md",
              "Abstractions" =>

docs/src/basics/conv.md

@@ -0,0 +1,3 @@
+# Graph convolutions
+
+Graph convolution can be classified into spectral-based graph convolution and spatial-based graph convolution. Spectral-based graph convolution, such as `GCNConv` and `ChebConv`, performs operation on features of *whole* graph at one time. Spatial-based graph convolution, such as `GraphConv` and `GATConv`, performs operation on features of *local* graph instead. Message-passing scheme is an abstraction for spatial-based graph convolutional layers. Any spatial-based graph convolutional layer can be implemented under the framework of message-passing scheme.

docs/src/basics/layers.md

@@ -2,7 +2,7 @@
 
 Building GNN is as simple as building neural network in Flux. The syntax here is the same as Flux. `Chain` is used to stack layers into a GNN. A simple example is shown here:
 
-```
+```julia
 model = Chain(GCNConv(adj_mat, feat=>h1),
               GCNConv(adj_mat, h1=>h2, relu))
 ```
@@ -21,7 +21,6 @@ When using GNN layers, the general guidelines are:
 * If you pass in a ``n \times d`` matrix of node features, and the layer maps node features ``\mathbb{R}^d \rightarrow \mathbb{R}^k`` then the output will be in matrix with dimensions ``n \times k``. The same ostensibly goes for edge features but as of now no layer type supports outputting new edge features.
 * If you pass in a `FeaturedGraph`, the output will be also be a `FeaturedGraph` with modified node (and/or edge) features. Add `node_feature` as the following entry in the Flux chain (or simply call `node_feature()` on the output) if you wish to subsequently convert them to matrix form.
 
-
-## Customize layers
+## Create custom layers
 
 Customizing your own GNN layers are the same as customizing layers in Flux. You may want to reference [Flux documentation](https://fluxml.ai/Flux.jl/stable/models/basics/#Building-Layers-1).

src/GeometricFlux.jl

@@ -5,6 +5,7 @@ using LinearAlgebra: Adjoint, norm, Transpose
 using Reexport
 
 using CUDA
+using ChainRulesCore: @non_differentiable
 using FillArrays: Fill
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
@@ -66,7 +67,6 @@ include("layers/pool.jl")
 include("models.jl")
 include("layers/misc.jl")
 
-include("cuda/msgpass.jl")
 include("cuda/conv.jl")
 
 using .Datasets

src/cuda/msgpass.jl

@@ -37,7 +37,9 @@ GCNConv(ch::Pair{Int,Int}, σ = identity; kwargs...) =
 @functor GCNConv
 
 function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
-    L̃ = normalized_laplacian(fg, eltype(x); selfloop=true)
+    L̃ = Zygote.ignore() do
+        GraphSignals.normalized_laplacian(fg, eltype(x); selfloop=true)
+    end
     l.σ.(l.weight * x * L̃ .+ l.bias)
 end
 
@@ -86,10 +88,12 @@ ChebConv(ch::Pair{Int,Int}, k::Int; kwargs...) =
 @functor ChebConv
 
 function (c::ChebConv)(fg::FeaturedGraph, X::AbstractMatrix{T}) where T
-    check_num_nodes(fg, X)
+    GraphSignals.check_num_nodes(fg, X)
     @assert size(X, 1) == size(c.weight, 2) "Input feature size must match input channel size."
 
-    L̃ = scaled_laplacian(fg, eltype(X))    
+    L̃ = Zygote.ignore() do
+        GraphSignals.scaled_laplacian(fg, eltype(X))
+    end
 
     Z_prev = X
     Z = X * L̃
@@ -156,12 +160,13 @@ 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)
-    check_num_nodes(fg, x)
-    _, x = propagate(gc, adjacency_list(fg), Fill(0.f0, 0, ne(fg)), x, +)
+    GraphSignals.check_num_nodes(fg, x)
+    _, x, _ = propagate(gc, graph(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
 
 function Base.show(io::IO, l::GraphConv)
     in_channel = size(l.weight1, ndims(l.weight1))
@@ -239,21 +244,21 @@ function apply_batch_message(gat::GATConv, i, js, X::AbstractMatrix)
     reshape(msgs, (n-1)*gat.heads, :)
 end
 
-update_batch_edge(gat::GATConv, adj, E::AbstractMatrix, X::AbstractMatrix, u) = update_batch_edge(gat, adj, X)
+function update_batch_edge(gat::GATConv, sg::SparseGraph, E::AbstractMatrix, X::AbstractMatrix, u)
+    @assert check_self_loops(sg) "a vertex must have self loop (receive a message from itself)."
+    mapreduce(i -> apply_batch_message(gat, i, neighbors(sg, i), X), hcat, 1:nv(sg))
+end
 
-function update_batch_edge(gat::GATConv, adj, X::AbstractMatrix)
-    n = size(adj, 1)
-    # a vertex must always receive a message from itself
-    Zygote.ignore() do
-        GraphLaplacians.add_self_loop!(adj, n)
+function check_self_loops(sg::SparseGraph)
+    for i in 1:nv(sg)
+        if !(i in GraphSignals.rowvalview(sg.S, i))
+            return false
+        end
     end
-    mapreduce(i -> apply_batch_message(gat, i, adj[i], X), hcat, 1:n)
+    return true
 end
 
-# The same as update function in batch manner
-update_batch_vertex(gat::GATConv, M::AbstractMatrix, X::AbstractMatrix, u) = update_batch_vertex(gat, M)
-
-function update_batch_vertex(gat::GATConv, M::AbstractMatrix)
+function update_batch_vertex(gat::GATConv, M::AbstractMatrix, X::AbstractMatrix, u)
     M = M .+ gat.bias
     if !gat.concat
         N = size(M, 2)
@@ -263,12 +268,13 @@ function update_batch_vertex(gat::GATConv, M::AbstractMatrix)
 end
 
 function (gat::GATConv)(fg::FeaturedGraph, X::AbstractMatrix)
-    check_num_nodes(fg, X)
-    _, X = propagate(gat, adjacency_list(fg), Fill(0.f0, 0, ne(fg)), X, +)
-    X
+    GraphSignals.check_num_nodes(fg, X)
+    _, X, _ = propagate(gat, graph(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
 
 function Base.show(io::IO, l::GATConv)
     in_channel = size(l.weight, ndims(l.weight))
@@ -319,10 +325,9 @@ update(ggc::GatedGraphConv, m::AbstractVector, x) = m
 
 
 function (ggc::GatedGraphConv)(fg::FeaturedGraph, H::AbstractMatrix{S}) where {T<:AbstractVector,S<:Real}
-    check_num_nodes(fg, H)
+    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."
-    adj = adjacency_list(fg)
     if m < ggc.out_ch
         Hpad = Zygote.ignore() do
             fill!(similar(H, S, ggc.out_ch - m, n), 0)
@@ -331,13 +336,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, adj, Fill(0.f0, 0, ne(fg)), M, +)
+        _, M = propagate(ggc, graph(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
 
 
 function Base.show(io::IO, l::GatedGraphConv)
@@ -374,12 +380,13 @@ message(ec::EdgeConv, x_i::AbstractVector, x_j::AbstractVector, e_ij) = ec.nn(vc
 update(ec::EdgeConv, m::AbstractVector, x) = m
 
 function (ec::EdgeConv)(fg::FeaturedGraph, X::AbstractMatrix)
-    check_num_nodes(fg, X)
-    _, X = propagate(ec, adjacency_list(fg), Fill(0.f0, 0, ne(fg)), X, ec.aggr)
+    GraphSignals.check_num_nodes(fg, X)
+    _, X, _ = propagate(ec, graph(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
 
 function Base.show(io::IO, l::EdgeConv)
     print(io, "EdgeConv(", l.nn)
@@ -425,12 +432,13 @@ update(g::GINConv, m::AbstractVector, x) = g.nn((1 + g.eps) * x + m)
 
 function (g::GINConv)(fg::FeaturedGraph, X::AbstractMatrix)
     gf = graph(fg)
-    GraphSignals.check_num_node(gf, X)
-    _, X = propagate(g, adjacency_list(gf), Fill(0.f0, 0, ne(gf)), X, +)
+    GraphSignals.check_num_nodes(gf, X)
+    _, X, _ = propagate(g, graph(fg), edge_feature(fg), X, global_feature(fg), +)
     X
 end
 
-(l::GINConv)(fg::FeaturedGraph) = FeaturedGraph(fg.graph, nf = l(fg, node_feature(fg)))
+(l::GINConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf = l(fg, node_feature(fg)))
+# (l::GINConv)(fg::FeaturedGraph) = propagate(l, fg, +)  # edge number check break this
 
 
 """
@@ -490,15 +498,16 @@ end
 update(c::CGConv, m::AbstractVector, x) = x + m
 
 function (c::CGConv)(fg::FeaturedGraph, X::AbstractMatrix, E::AbstractMatrix)
-    check_num_nodes(fg, X)
-    check_num_edges(fg, E)
-    _, Y = propagate(c, adjacency_list(fg), E, X, +)
+    GraphSignals.check_num_nodes(fg, X)
+    GraphSignals.check_num_edges(fg, E)
+    _, Y, _ = propagate(c, graph(fg), E, X, global_feature(fg), +)
     Y
 end
 
-(l::CGConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf=l(fg, node_feature(fg),
-                                                        edge_feature(fg)),
+(l::CGConv)(fg::FeaturedGraph) = 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)(X::AbstractMatrix, E::AbstractMatrix) = l(l.fg, X, E)
 

src/layers/conv.jl

@@ -16,52 +16,32 @@ 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, adj, E, V, u)
-    n = size(adj, 1)
-    edge_idx = edge_index_table(adj)
-    mapreduce(i -> apply_batch_message(gn, i, adj[i], edge_idx, E, V, u), hcat, 1:n)
-end
+@inline update_batch_edge(gn::GraphNet, sg::SparseGraph, E, V, u) =
+    mapreduce(i -> apply_batch_message(gn, sg, i, neighbors(sg, i), E, V, u), hcat, vertices(sg))
 
-@inline apply_batch_message(gn::GraphNet, i, js, edge_idx, E, V, u) =
-    mapreduce(j -> update_edge(gn, _view(E, edge_idx[(i,j)]), _view(V, i), _view(V, j), u), hcat, js)
+@inline apply_batch_message(gn::GraphNet, sg::SparseGraph, i, js, E, V, u) =
+    mapreduce(j -> update_edge(gn, _view(E, edge_index(sg, i, j)), _view(V, i), _view(V, j), u), hcat, js)
 
 @inline update_batch_vertex(gn::GraphNet, Ē, V, u) =
     mapreduce(i -> update_vertex(gn, _view(Ē, i), _view(V, i), u), hcat, 1:size(V,2))
 
-@inline function aggregate_neighbors(gn::GraphNet, aggr, E, accu_edge)
-    @assert !iszero(accu_edge) "accumulated edge must not be zero."
-    cluster = generate_cluster(E, accu_edge)
-    NNlib.scatter(aggr, E, cluster)
-end
-
-@inline function aggregate_neighbors(gn::GraphNet, aggr::Nothing, E, accu_edge)
-    @nospecialize E accu_edge
-    return nothing
-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_edges(gn::GraphNet, aggr, E) = aggregate(aggr, E)
-
-@inline function aggregate_edges(gn::GraphNet, aggr::Nothing, E)
-    @nospecialize E
-    return nothing
-end
+@inline aggregate_edges(gn::GraphNet, aggr::Nothing, @nospecialize E) = nothing
 
 @inline aggregate_vertices(gn::GraphNet, aggr, V) = aggregate(aggr, V)
-
-@inline function aggregate_vertices(gn::GraphNet, aggr::Nothing, V)
-    @nospecialize V
-    return nothing
-end
+@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, adjacency_list(fg), fg.ef, fg.nf, fg.gf, naggr, eaggr, vaggr)
+    E, V, u = propagate(gn, graph(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, adj::AbstractVector{S}, E::R, V::Q, u::P,
-                   naggr=nothing, eaggr=nothing, vaggr=nothing) where {S<:AbstractVector,R,Q,P}
-    E = update_batch_edge(gn, adj, E, V, u)
-    Ē = aggregate_neighbors(gn, naggr, E, accumulated_edges(adj))
+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)
     ē = aggregate_edges(gn, eaggr, E)
     v̄ = aggregate_vertices(gn, vaggr, V)

src/layers/gn.jl

@@ -22,6 +22,8 @@ specialize this method with custom behavior.
 
 See also [`update`](@ref).
 """
+function message end
+
 @inline message(mp::MessagePassing, x_i, x_j, e_ij) = x_j
 @inline message(mp::MessagePassing, i::Integer, j::Integer, x_i, x_j, e_ij) = x_j
 
@@ -45,21 +47,10 @@ specialize this method with custom behavior.
 
 See also [`message`](@ref).
 """
+function update end
+
 @inline update(mp::MessagePassing, m, x) = m
 @inline update(mp::MessagePassing, i::Integer, m, x) = m
 
-@inline apply_batch_message(mp::MessagePassing, i, js, edge_idx, E::AbstractMatrix, X::AbstractMatrix, u) =
-    mapreduce(j -> GeometricFlux.message(mp, _view(X, i), _view(X, j), _view(E, edge_idx[(i,j)])), hcat, js)
-
-@inline update_batch_vertex(mp::MessagePassing, M::AbstractMatrix, X::AbstractMatrix, u) = 
-    mapreduce(i -> GeometricFlux.update(mp, _view(M, i), _view(X, i)), hcat, 1:size(X,2))
-
-function propagate(mp::MessagePassing, fg::FeaturedGraph, aggr=+)
-    E, X = propagate(mp, adjacency_list(fg), fg.ef, fg.nf, aggr)
-    FeaturedGraph(fg, nf=X, ef=E, gf=Fill(0.f0, 0))
-end
-
-function propagate(mp::MessagePassing, adj::AbstractVector{S}, E::R, X::Q, aggr) where {S<:AbstractVector,R,Q}
-    E, X, u = propagate(mp, adj, E, X, Fill(0.f0, 0), aggr, nothing, nothing)
-    E, X
-end
+@inline update_edge(mp::MessagePassing, e, vi, vj, u) = GeometricFlux.message(mp, vi, vj, e)
+@inline update_vertex(mp::MessagePassing, ē, vi, u) = GeometricFlux.update(mp, ē, vi)