support edge_weight in GCNConv

Author: CarloLucibello

src/GNNGraphs/query.jl

@@ -95,18 +95,23 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g.g
     return dir == :out ? A : A'
 end
 
-function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=true)
-    s, t = edge_index(g)
-
+function _get_edge_weight(g, edge_weight)
     if edge_weight === true
-        edge_weight = get_edge_weight(g)
+        ew = get_edge_weight(g)
     elseif (edge_weight === false) || (edge_weight === nothing)
-        edge_weight = nothing 
+        ew = nothing 
     elseif edge_weight isa AbstractVector
-        edge_weight = edge_weight 
-    else 
+        ew = edge_weight 
+    else
         error("Invalid edge_weight argument.")
     end
+    return ew
+end
+
+function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=true)
+    s, t = edge_index(g)
+
+    edge_weight = _get_edge_weight(g, edge_weight)
     edge_weight = isnothing(edge_weight) ? eltype(s)(1) : edge_weight
 
     T = isnothing(T) ? eltype(edge_weight) : T

src/layers/conv.jl

@@ -1,57 +1,81 @@
 @doc raw"""
-    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform, add_self_loops=true)
+    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform, add_self_loops=true, edge_weight=true)
 
 Graph convolutional layer from paper [Semi-supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).
 
 Performs the operation
 ```math
-\mathbf{x}'_i = \sum_{j\in N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
+\mathbf{x}'_i = \sum_{j\in N(i)} a_{ij} W \mathbf{x}_j
 ```
-where ``c_{ij} = \sqrt{|N(i)||N(j)|}``.
+where ``a_{ij} = 1 / \sqrt{|N(i)||N(j)|}`` is a normalization factor computed from the node degrees. 
 
-The input to the layer is a node feature array `X` 
-of size `(num_features, num_nodes)`.
+If the input graph has weighted edges and `edge_weight=true`, than ``c_{ij}`` will be computed as
+```math
+a_{ij} = \frac{e_{j\to i}}{\sqrt{\sum_{j \in N(i)}  e_{j\to i}} \sqrt{\sum_{i \in N(j)}  e_{i\to j}}}
+```
+
+The input to the layer is a node feature array `X` of size `(num_features, num_nodes)`
+and optionally an edge weight vector.
 
 # Arguments
 
 - `in`: Number of input features.
 - `out`: Number of output features.
-- `σ`: Activation function.
-- `bias`: Add learnable bias.
-- `init`: Weights' initializer.
-- `add_self_loops`: Add self loops to the graph before performing the convolution.
+- `σ`: Activation function. Default `identity`.
+- `bias`: Add learnable bias. Default `true`.
+- `init`: Weights' initializer. Default `glorot_uniform`.
+- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
+- `edge_weight`. If `true`, consider the edge weights in the input graph (if available).
+                 Not compatible with `add_self_loops=true` at the moment. Default `true`.
 """
 struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
     bias::B
     σ::F
     add_self_loops::Bool
+    edge_weight::Bool
 end
 
 @functor GCNConv
 
 function GCNConv(ch::Pair{Int,Int}, σ=identity;
-                 init=glorot_uniform, bias::Bool=true,
-                 add_self_loops=true)
+                 init=glorot_uniform, 
+                 bias::Bool=true,
+                 add_self_loops=true,
+                 edge_weight=false)
     in, out = ch
     W = init(out, in)
     b = bias ? Flux.create_bias(W, true, out) : false
-    GCNConv(W, b, σ, add_self_loops)
+    GCNConv(W, b, σ, add_self_loops, edge_weight)
 end
 
-function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}) where T
+function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix)
+    # Extract edge_weight from g if available and l.edge_weight == false,
+    # otherwise return nothing.
+    edge_weight = GNNGraphs._get_edge_weight(g, l.edge_weight) # vector or nothing
+    return l(g, x, edge_weight)
+end
+
+function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}, edge_weight::EW) where 
+    {T, EW<:Union{Nothing,AbstractVector}}
+    
     if l.add_self_loops
+        @assert edge_weight === nothing
         g = add_self_loops(g)
     end
     Dout, Din = size(l.weight)
     if Dout < Din
         x = l.weight * x
     end
     # @assert all(>(0), degree(g, T, dir=:in))
-    c = 1 ./ sqrt.(degree(g, T, dir=:in))
-    x = x .* c'
-    x = propagate(copy_xj, g, +, xj=x)
+    c = 1 ./ sqrt.(degree(g, T; dir=:in, edge_weight))
     x = x .* c'
+    if edge_weight === nothing
+        x = propagate(copy_xj, g, +, xj=x)
+    else        
+        x = propagate(e_mul_xj, g, +, xj=x, e=edge_weight)
+    end
+        x = x .* c'
     if Dout >= Din
         x = l.weight * x
     end

src/msgpass.jl

@@ -153,6 +153,18 @@ copy_xi(xi, xj, e) = xi
 """
 xi_dot_xj(xi, xj, e) = sum(xi .* xj, dims=1)
 
+"""
+    e_mul_xj(xi, xj, e) = reshape(e, (...)) .* xj
+
+Reshape `e` into broadcast compatible shape with `xj`
+(by prepending singleton dimensions) then perform
+broadcasted multiplication.
+"""
+function e_mul_xj(xi, xj::AbstractArray{Tj,Nj}, e::AbstractArray{Te,Ne}) where {Tj,Te, Nj, Ne}
+    @assert Ne <= Nj
+    e = reshape(e, ntuple(_ -> 1, Nj-Ne)..., size(e)...)
+    return e .* xj
+end
 
 function propagate(::typeof(copy_xj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
     A = adjacency_matrix(g)

test/runtests.jl

@@ -21,18 +21,18 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets
 include("test_utils.jl")
 
 tests = [
-    "GNNGraphs/gnngraph",
-    "GNNGraphs/transform",
-    "GNNGraphs/operators",
-    "GNNGraphs/generate",
-    "GNNGraphs/query",
-    "utils",
-    "msgpass",
-    "layers/basic",
+    # "GNNGraphs/gnngraph",
+    # "GNNGraphs/transform",
+    # "GNNGraphs/operators",
+    # "GNNGraphs/generate",
+    # "GNNGraphs/query",
+    # "utils",
+    # "msgpass",
+    # "layers/basic",
     "layers/conv",
-    "layers/pool",
-    "examples/node_classification_cora",
-    "deprecations",
+    # "layers/pool",
+    # "examples/node_classification_cora",
+    # "deprecations",
 ]
 
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")