Merge pull request #36 from CarloLucibello/cl/ecconv

add NNConv and FiniteDifferences testing

Author: CarloLucibello

Project.toml

@@ -1,7 +1,7 @@
 name = "GraphNeuralNetworks"
 uuid = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 authors = ["Carlo Lucibello and contributors"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -25,16 +25,19 @@ ChainRulesCore = "1"
 DataStructures = "0.18"
 Flux = "0.12"
 KrylovKit = "0.5"
-LearnBase = "0.5"
+LearnBase = "0.4, 0.5"
 LightGraphs = "1.3"
 MacroTools = "0.5"
 NNlib = "0.7"
 NNlibCUDA = "0.1"
 julia = "1.6"
 
 [extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Zygote"]
+test = ["Test", "Adapt", "Zygote", "FiniteDifferences", "ChainRulesTestUtils"]

src/GraphNeuralNetworks.jl

@@ -38,13 +38,14 @@ export
     GNNChain,
 
     # layers/conv
-    GCNConv,
     ChebConv,
-    GraphConv,
+    EdgeConv,
     GATConv,
     GatedGraphConv,
-    EdgeConv,
+    GCNConv,
     GINConv,
+    GraphConv,
+    NNConv,
 
     # layers/pool
     GlobalPool,

src/gnngraph.jl

@@ -151,6 +151,10 @@ GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)
 function GNNGraph(g::AbstractGraph; kws...)
     s = LightGraphs.src.(LightGraphs.edges(g))
     t = LightGraphs.dst.(LightGraphs.edges(g))
+    if !LightGraphs.is_directed(g) 
+        # add reverse edges since GNNGraph are directed
+        s, t = [s; t], [t; s]
+    end
     GNNGraph((s, t); num_nodes = LightGraphs.nv(g), kws...)
 end
 

src/layers/basic.jl

@@ -7,7 +7,8 @@ See also [`GNNChain`](@ref).
 """
 abstract type GNNLayer end
 
-#TODO extend to store also edge and global features
+# Forward pass with graph-only input.
+# To be specialized by layers also needing edge features as input (e.g. NNConv). 
 (l::GNNLayer)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g)))
 
 """

src/layers/conv.jl

@@ -32,7 +32,7 @@ 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)
+    b = bias ? Flux.create_bias(W, true, out) : false
     GCNConv(W, b, σ)
 end
 
@@ -105,7 +105,7 @@ function ChebConv(ch::Pair{Int,Int}, k::Int;
                   init=glorot_uniform, bias::Bool=true)
     in, out = ch
     W = init(out, in, k)
-    b = Flux.create_bias(W, bias, out)
+    b = bias ? Flux.create_bias(W, true, out) : false
     ChebConv(W, b, k)
 end
 
@@ -172,7 +172,7 @@ function GraphConv(ch::Pair{Int,Int}, σ=identity, aggr=+;
     in, out = ch
     W1 = init(out, in)
     W2 = init(out, in)
-    b = Flux.create_bias(W1, bias, out)
+    b = bias ? Flux.create_bias(W1, true, out) : false
     GraphConv(W1, W2, b, σ, aggr)
 end
 
@@ -196,7 +196,7 @@ end
 
 
 @doc raw"""
-    GATConv(in => out, , σ=identity;
+    GATConv(in => out, σ=identity;
             heads=1,
             concat=true,
             init=glorot_uniform    
@@ -224,7 +224,7 @@ with ``z_i`` a normalization factor.
 - `concat`: Concatenate layer output or not. If not, layer output is averaged over the heads.
 - `negative_slope::Real`: Keyword argument, the parameter of LeakyReLU.
 """
-struct GATConv{T, A<:AbstractMatrix{T}, B} <: GNNLayer
+struct GATConv{T, A<:AbstractMatrix, B} <: GNNLayer
     weight::A
     bias::B
     a::A
@@ -239,12 +239,13 @@ end
 Flux.trainable(l::GATConv) = (l.weight, l.bias, l.a)
 
 function GATConv(ch::Pair{Int,Int}, σ=identity;
-                 heads::Int=1, concat::Bool=true, negative_slope=0.2f0,
+                 heads::Int=1, concat::Bool=true, negative_slope=0.2,
                  init=glorot_uniform, bias::Bool=true)
     in, out = ch             
     W = init(out*heads, in)
-    b = Flux.create_bias(W, bias, out*heads)
+    b = bias ? Flux.create_bias(W, true, out*heads) : false
     a = init(2*out, heads)
+    negative_slope = convert(eltype(W), negative_slope)
     GATConv(W, b, a, σ, negative_slope, ch, heads, concat)
 end
 
@@ -356,20 +357,20 @@ end
 
 
 @doc raw"""
-    EdgeConv(f; aggr=max)
+    EdgeConv(nn; aggr=max)
 
 Edge convolutional layer from paper [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829).
 
 Performs the operation
 ```math
-\mathbf{x}_i' = \square_{j \in N(i)} f(\mathbf{x}_i || \mathbf{x}_j - \mathbf{x}_i)
+\mathbf{x}_i' = \square_{j \in N(i)} nn(\mathbf{x}_i || \mathbf{x}_j - \mathbf{x}_i)
 ```
 
-where `f` typically denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.
+where `nn` generally denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.
 
 # Arguments
 
-- `f`: A (possibly learnable) function acting on edge features. 
+- `nn`: A (possibly learnable) function acting on edge features. 
 - `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
 """
 struct EdgeConv <: GNNLayer
@@ -405,9 +406,9 @@ Graph Isomorphism convolutional layer from paper [How Powerful are Graph Neural
 
 
 ```math
-\mathbf{x}_i' = f\left((1 + \epsilon) \mathbf{x}_i + \sum_{j \in N(i)} \mathbf{x}_j \right)
+\mathbf{x}_i' = f_\Theta\left((1 + \epsilon) \mathbf{x}_i + \sum_{j \in N(i)} \mathbf{x}_j \right)
 ```
-where `f` typically denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.
+where ``f_\Theta`` typically denotes a learnable function, e.g. a linear layer or a multi-layer perceptron.
 
 # Arguments
 
@@ -434,3 +435,77 @@ function (l::GINConv)(g::GNNGraph, X::AbstractMatrix)
     X, _ = propagate(l, g, +, X)
     X
 end
+
+
+@doc raw"""
+    NNConv(in => out, f, σ=identity; aggr=+, bias=true, init=glorot_uniform)
+
+The continuous kernel-based convolutional operator from the 
+[Neural Message Passing for Quantum Chemistry](https://arxiv.org/abs/1704.01212) paper. 
+This convolution is also known as the edge-conditioned convolution from the 
+[Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) paper.
+
+Performs the operation
+
+```math
+\mathbf{x}_i' = W \mathbf{x}_i + \square_{j \in N(i)} f_\Theta(\mathbf{e}_{j\to i})\,\mathbf{x}_j
+```
+
+where ``f_\Theta``  denotes a learnable function (e.g. a linear layer or a multi-layer perceptron).
+Given an input of batched edge features `e` of size `(num_edge_features, num_edges)`, 
+the function `f` will return an batched matrices array whose size is `(out, in, num_edges)`.
+For convenience, also functions returning a single `(out*in, num_edges)` matrix are allowed.
+
+# Arguments
+
+- `in`: The dimension of input features.
+- `out`: The dimension of output features.
+- `f`: A (possibly learnable) function acting on edge features.
+- `aggr`: Aggregation operator for the incoming messages (e.g. `+`, `*`, `max`, `min`, and `mean`).
+- `σ`: Activation function.
+- `bias`: Add learnable bias.
+- `init`: Weights' initializer.
+"""
+struct NNConv <: GNNLayer
+    weight
+    bias
+    nn
+    σ
+    aggr
+end
+
+@functor NNConv
+
+function NNConv(ch::Pair{Int,Int}, nn, σ=identity; aggr=+, bias=true, init=glorot_uniform)
+    in, out = ch
+    W = init(out, in)
+    b = bias ? Flux.create_bias(W, true, out) : false
+    return NNConv(W, b, nn, σ, aggr)
+end
+
+function compute_message(l::NNConv, x_i, x_j, e_ij) 
+    nin, nedges = size(x_i)
+    W = reshape(l.nn(e_ij), (:, nin, nedges))
+    x_j = reshape(x_j, (nin, 1, nedges)) # needed by batched_mul
+    m = NNlib.batched_mul(W, x_j)
+    return reshape(m, :, nedges)
+end
+
+function update_node(l::NNConv, m, x)
+    l.σ.(l.weight*x .+ m .+ l.bias)
+end
+
+function (l::NNConv)(g::GNNGraph, x::AbstractMatrix, e)
+    check_num_nodes(g, x)
+    x, _ = propagate(l, g, l.aggr, x, e)
+    return x
+end
+
+(l::NNConv)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g), edge_features(g)))
+
+function Base.show(io::IO, l::NNConv)
+    out, in = size(l.weight)
+    print(io, "NNConv( $in => $out")
+    print(io, ", aggr=", l.aggr)
+    print(io, ")")
+end

test/cuda/layers/conv.jl

@@ -84,6 +84,19 @@
         end
     end
 
+    @testset "LightGraphs constructor" begin
+        lg = random_regular_graph(10, 4)
+        @test !LightGraphs.is_directed(lg)
+        g = GNNGraph(lg)
+        @test g.num_edges == 2*ne(lg) # g in undirected
+        @test LightGraphs.is_directed(g)
+        for e in LightGraphs.edges(lg)
+            i, j = src(e), dst(e)
+            @test has_edge(g, i, j)
+            @test has_edge(g, j, i)            
+        end
+    end
+
     @testset "add self-loops" begin
         A = [1  1  0  0
              0  0  1  0
@@ -174,9 +187,9 @@
     @testset "LearnBase and DataLoader compat" begin
         n, m, num_graphs = 10, 30, 50
         X = rand(10, n)
-        E = rand(10, m)
+        E = rand(10, 2m)
         U = rand(10, 1)
-        g = Flux.batch([GNNGraph(erdos_renyi(10, 30), ndata=rand(10, n), edata=rand(10, m), gdata=rand(10, 1)) 
+        g = Flux.batch([GNNGraph(erdos_renyi(n, m), ndata=X, edata=E, gdata=U) 
                         for _ in 1:num_graphs])
 
         @test LearnBase.getobs(g, 3) == getgraph(g, 3)[1]