Merge pull request #119 from yuehhua/refactor

Refactor

Author: yuehhua

src/GeometricFlux.jl

@@ -5,6 +5,7 @@ using SparseArrays: SparseMatrixCSC
 using LinearAlgebra: Adjoint, norm
 using Reexport
 
+using CUDA
 using FillArrays: Fill
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell
@@ -105,7 +106,6 @@ using .Datasets
 
 function __init__()
     @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin
-        using CUDA
         include("cuda/msgpass.jl")
         include("cuda/conv.jl")
         include("cuda/pool.jl")

src/cuda/msgpass.jl

@@ -9,16 +9,13 @@ end
 end
 
 @inline function update_batch_edge(mp::T, adj, E::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
+    n = size(adj, 1)
     edge_idx = edge_index_table(adj)
-    E_ = Vector[]
-    for (i, js) = enumerate(adj)
-        for j = js
-            k = edge_idx[(i,j)]
-            m = message(mp, get_feature(X, i), get_feature(X, j), get_feature(E, k))
-            push!(E_, m)
-        end
-    end
-    hcat(E_...)
+    hcat([apply_batch_message(mp, i, adj[i], edge_idx, E, X) for i in 1:n]...)
+end
+
+@inline function apply_batch_message(mp::T, i, js, edge_idx, E::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
+    hcat([message(mp, get_feature(X, i), get_feature(X, j), get_feature(E, edge_idx[(i,j)])) for j = js]...)
 end
 
 @inline function update_batch_vertex(mp::T, M::AbstractMatrix, X::CuMatrix) where {T<:MessagePassing}
@@ -32,10 +29,6 @@ end
 end
 
 @inline function update_batch_vertex(mp::T, M::CuMatrix, X::CuMatrix) where {T<:MessagePassing}
-    X_ = Vector[]
-    for i = 1:size(X,2)
-        x = update(mp, get_feature(M, i), get_feature(X, i))
-        push!(X_, x)
-    end
+    X_ = [update(mp, get_feature(M, i), get_feature(X, i)) for i = 1:size(X,2)]
     hcat(X_...)
 end

src/graph/index.jl

@@ -53,20 +53,31 @@ end
 Zygote.@nograd vertex_pair_table
 
 """
-    edge_index_table(adj[, num_E])
+    edge_index_table(adj[, directed])
 
 Generate a mapping from vertex pair (i, j) to edge index. The edge indecies are determined by
 the sorted vertex indecies.
 """
-function edge_index_table(adj::AbstractVector{<:AbstractVector{<:Integer}},
-                          num_E=sum(map(length, adj)))
+function edge_index_table(adj::AbstractVector{<:AbstractVector{<:Integer}}, directed::Bool=is_directed(adj))
     table = Dict{Tuple{UInt32,UInt32},UInt64}()
     e = one(UInt64)
-    for (i, js) = enumerate(adj)
-        js = sort(js)
-        for j = js
-            table[(i, j)] = e
-            e += one(UInt64)
+    if directed
+        for (i, js) = enumerate(adj)
+            js = sort(js)
+            for j = js
+                table[(i, j)] = e
+                e += one(UInt64)
+            end
+        end
+    else
+        for (i, js) = enumerate(adj)
+            js = sort(js)
+            js = js[i .≤ js]
+            for j = js
+                table[(i, j)] = e
+                table[(j, i)] = e
+                e += one(UInt64)
+            end
         end
     end
     table
@@ -80,6 +91,8 @@ function edge_index_table(vpair::AbstractVector{<:Tuple})
     table
 end
 
+edge_index_table(fg::FeaturedGraph) = edge_index_table(fg.graph, fg.directed)
+
 Zygote.@nograd edge_index_table
 
 function transform(X::AbstractArray, vpair::AbstractVector{<:Tuple}, num_V)

src/layers/gn.jl

@@ -10,11 +10,11 @@ abstract type GraphNet end
 @inline function update_batch_edge(gn::T, adj, E, V, u) where {T<:GraphNet}
     n = size(adj, 1)
     edge_idx = edge_index_table(adj)
-    E_ = [_apply_batch_message(gn, i, adj[i], edge_idx, E, V, u) for i in 1:n]
+    E_ = [apply_batch_message(gn, i, adj[i], edge_idx, E, V, u) for i in 1:n]
     hcat(E_...)
 end
 
-@inline function _apply_batch_message(gn::T, i, js, edge_idx, E, V, u) where {T<:GraphNet}
+@inline function apply_batch_message(gn::T, i, js, edge_idx, E, V, u) where {T<:GraphNet}
     E_ = [update_edge(gn, get_feature(E, edge_idx[(i,j)]), get_feature(V, i), get_feature(V, j), u) for j = js]
     hcat(E_...)
 end

src/utils.jl

@@ -1,30 +1,3 @@
-## Indexing
-
-function range_indecies(idx::Tuple)
-    x = Vector{Any}(undef, length(idx))
-    for (i,n) in enumerate(idx)
-        x[i] = 1:n
-    end
-    x
-end
-
-replace_last_index!(idx::Vector, x) = (idx[end] = x; idx)
-
-function assign!(A::AbstractArray, B::AbstractArray{T,N}; last_dim=1:size(B,N)) where {T,N}
-    A_dims, B_dims = size(A), size(B)
-    @assert A_dims[1:end-1] == B_dims[1:end-1] "Inconsistent dimensions with $(A_dims[1:end-1]) and $(B_dims[1:end-1])"
-    A_dims = replace_last_index!(range_indecies(A_dims), last_dim)
-    B_dims = range_indecies(B_dims)
-    A_idxs = CartesianIndices(Tuple(A_dims))
-    B_idxs = CartesianIndices(Tuple(B_dims))
-    for (Aidx, Bidx) = zip(A_idxs, B_idxs)
-        A[Aidx] = B[Bidx]
-    end
-    A
-end
-
-
-
 ## Top-k pooling
 
 function topk_index(y::AbstractVector, k::Integer)
@@ -38,8 +11,9 @@ topk_index(y::Adjoint, k::Integer) = topk_index(y', k)
 
 ## Get feature with defaults
 
-get_feature(::Nothing, i::Integer) = zeros(0)
-get_feature(A::AbstractMatrix, i::Integer) = (i ≤ size(A,2)) ? view(A, :, i) : zeros(0)
+get_feature(::Nothing, i) = nothing
+get_feature(A::Fill{T,2,Axes}, i::Integer) where {T,Axes} = view(A, :, 1)
+get_feature(A::AbstractMatrix, i::Integer) = view(A, :, i)
 
 """
     bypass_graph(nf_func, ef_func, gf_func)

test/cuda/msgpass.jl

@@ -1,6 +1,7 @@
 in_channel = 10
 out_channel = 5
 N = 6
+T = Float32
 adj = [0. 1. 0. 0. 0. 0.;
        1. 0. 0. 1. 1. 1.;
        0. 0. 0. 0. 0. 1.;
@@ -18,8 +19,9 @@ NewCudaLayer(m, n) = NewCudaLayer(randn(m,n))
 GeometricFlux.message(n::NewCudaLayer, x_i, x_j, e_ij) = n.weight * x_j
 GeometricFlux.update(::NewCudaLayer, m, x) = m
 
-X = rand(Float32, in_channel, N) |> gpu
+X = rand(T, in_channel, N) |> gpu
 fg = FeaturedGraph(adj, X)
+fg.ef = Fill(zero(T), 0, 2num_E)
 l = NewCudaLayer(out_channel, in_channel) |> gpu
 
 @testset "cuda/msgpass" begin

test/layers/msgpass.jl

@@ -2,6 +2,7 @@ in_channel = 10
 out_channel = 5
 num_V = 6
 num_E = 7
+T = Float32
 adj = [0. 1. 0. 0. 0. 0.;
        1. 0. 0. 1. 1. 1.;
        0. 0. 0. 0. 0. 1.;
@@ -16,8 +17,9 @@ NewLayer(m, n) = NewLayer(randn(m,n))
 
 (l::NewLayer)(fg) = propagate(l, fg, :add)
 
-X = Array(reshape(1:num_V*in_channel, in_channel, num_V))
+X = Array{T}(reshape(1:num_V*in_channel, in_channel, num_V))
 fg = FeaturedGraph(adj, X)
+fg.ef = Fill(zero(T), 0, 2num_E)
 
 l = NewLayer(out_channel, in_channel)
 

test/runtests.jl

@@ -2,6 +2,7 @@ using GeometricFlux
 using GeometricFlux.Datasets
 using Flux
 using Flux: @functor
+using FillArrays
 using GraphSignals
 using StaticArrays: @MMatrix, @MArray
 using LightGraphs: SimpleGraph, SimpleDiGraph, add_edge!, nv, ne