Merge pull request #122 from CarloLucibello/tclements/master

faster batching

Author: CarloLucibello

src/GNNGraphs/transform.jl

@@ -270,12 +270,12 @@ function SparseArrays.blockdiag(g1::GNNGraph, g2::GNNGraph)
         t = vcat(t1, nv1 .+ t2)
         w = cat_features(get_edge_weight(g1), get_edge_weight(g2))
         graph = (s, t, w)
-        ind1 = isnothing(g1.graph_indicator) ? ones_like(s1, Int, nv1) : g1.graph_indicator 
-        ind2 = isnothing(g2.graph_indicator) ? ones_like(s2, Int, nv2) : g2.graph_indicator     
+        ind1 = isnothing(g1.graph_indicator) ? ones_like(s1, nv1) : g1.graph_indicator 
+        ind2 = isnothing(g2.graph_indicator) ? ones_like(s2, nv2) : g2.graph_indicator     
     elseif g1.graph isa ADJMAT_T        
         graph = blockdiag(g1.graph, g2.graph)
-        ind1 = isnothing(g1.graph_indicator) ? ones_like(graph, Int, nv1) : g1.graph_indicator 
-        ind2 = isnothing(g2.graph_indicator) ? ones_like(graph, Int, nv2) : g2.graph_indicator     
+        ind1 = isnothing(g1.graph_indicator) ? ones_like(graph, nv1) : g1.graph_indicator 
+        ind2 = isnothing(g2.graph_indicator) ? ones_like(graph, nv2) : g2.graph_indicator     
     end
     graph_indicator = vcat(ind1, g1.num_graphs .+ ind2)
 
@@ -358,8 +358,37 @@ julia> g12.ndata.x
  1.0  1.0  1.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 ```
 """
-Flux.batch(gs::Vector{<:GNNGraph}) = blockdiag(gs...)
+Flux.batch(gs::AbstractVector{<:GNNGraph}) = blockdiag(gs...)
+
+function Flux.batch(gs::AbstractVector{<:GNNGraph{T}}) where T<:COO_T
+    v_num_nodes = [g.num_nodes for g in gs]
+    edge_indices = [edge_index(g) for g in gs]
+    nodesum = cumsum([0; v_num_nodes])[1:end-1]
+    s = cat_features([ei[1] .+ nodesum[ii] for (ii, ei) in enumerate(edge_indices)])
+    t = cat_features([ei[2] .+ nodesum[ii] for (ii, ei) in enumerate(edge_indices)])
+    w = cat_features([get_edge_weight(g) for g in gs])
+    graph = (s, t, w)
+
+    function materialize_graph_indicator(g)
+        g.graph_indicator === nothing ? ones_like(s, g.num_nodes) : g.graph_indicator
+    end
 
+    v_gi = materialize_graph_indicator.(gs)
+    v_num_graphs = [g.num_graphs for g in gs]
+    graphsum = cumsum([0; v_num_graphs])[1:end-1]
+    v_gi = [ng .+ gi  for (ng, gi) in zip(graphsum, v_gi)]
+    graph_indicator = cat_features(v_gi)
+    
+    GNNGraph(graph,
+        sum(v_num_nodes), 
+        sum([g.num_edges for g in gs]),
+        sum(v_num_graphs),
+        graph_indicator,
+        cat_features([g.ndata for g in gs]),
+        cat_features([g.edata for g in gs]),
+        cat_features([g.gdata for g in gs]),
+    )
+end
 
 """
     unbatch(g::GNNGraph)

src/GNNGraphs/utils.jl

@@ -20,13 +20,33 @@ cat_features(x1::Union{Number, AbstractVector}, x2::Union{Number, AbstractVector
 
 # workaround for issue #98 #104
 cat_features(x1::NamedTuple{(), Tuple{}}, x2::NamedTuple{(), Tuple{}}) = (;)
+cat_features(xs::AbstractVector{NamedTuple{(), Tuple{}}}) = (;)
 
 function cat_features(x1::NamedTuple, x2::NamedTuple)
     sort(collect(keys(x1))) == sort(collect(keys(x2))) || @error "cannot concatenate feature data with different keys"
 
     NamedTuple(k => cat_features(getfield(x1,k), getfield(x2,k)) for k in keys(x1))
 end
 
+function cat_features(xs::AbstractVector{<:AbstractArray{T,N}}) where {T<:Number, N}
+   cat(xs...; dims=N) 
+end
+
+cat_features(xs::AbstractVector{Nothing}) = nothing
+cat_features(xs::AbstractVector{<:Number}) = xs
+
+function cat_features(xs::AbstractVector{<:NamedTuple})
+    symbols = [sort(collect(keys(x))) for x in xs]
+    all(y -> y==symbols[1], symbols) || @error "cannot concatenate feature data with different keys"
+    length(xs) == 1 && return xs[1] 
+
+    # concatenate 
+    syms = symbols[1]
+    NamedTuple(
+        k => cat_features([x[k] for x in xs]) for (ii,k) in enumerate(syms)
+    )
+end
+
 # Turns generic type into named tuple
 normalize_graphdata(data::Nothing; kws...) = NamedTuple()
 
@@ -70,9 +90,10 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
     return data
 end
 
-ones_like(x::AbstractArray, T=eltype(x), sz=size(x)) = fill!(similar(x, T, sz), 1)
-ones_like(x::SparseMatrixCSC, T=eltype(x), sz=size(x)) = ones(T, sz)
-ones_like(x::CUMAT_T, T=eltype(x), sz=size(x)) = CUDA.ones(T, sz)
+ones_like(x::AbstractArray, T::Type, sz=size(x)) = fill!(similar(x, T, sz), 1)
+ones_like(x::SparseMatrixCSC, T::Type, sz=size(x)) = ones(T, sz)
+ones_like(x::CUMAT_T, T::Type, sz=size(x)) = CUDA.ones(T, sz)
+ones_like(x, sz=size(x)) = ones_like(x, eltype(x), sz)
 
 numnonzeros(a::AbstractSparseMatrix) = nnz(a)
 numnonzeros(a::AbstractMatrix) = count(!=(0), a)

test/GNNGraphs/transform.jl

@@ -18,13 +18,13 @@
     end
 
     @testset "batch"  begin
-        #TODO add graph_type=GRAPH_T
-        g1 = GNNGraph(random_regular_graph(10,2), ndata=rand(16,10))
-        g2 = GNNGraph(random_regular_graph(4,2), ndata=rand(16,4))
-        g3 = GNNGraph(random_regular_graph(7,2), ndata=rand(16,7))
+        g1 = GNNGraph(random_regular_graph(10,2), ndata=rand(16,10), graph_type=GRAPH_T)
+        g2 = GNNGraph(random_regular_graph(4,2), ndata=rand(16,4), graph_type=GRAPH_T)
+        g3 = GNNGraph(random_regular_graph(7,2), ndata=rand(16,7), graph_type=GRAPH_T)
 
         g12 = Flux.batch([g1, g2])
         g12b = blockdiag(g1, g2)
+        @test g12 == g12b
 
         g123 = Flux.batch([g1, g2, g3])
         @test g123.graph_indicator == [fill(1, 10); fill(2, 4); fill(3, 7)]
@@ -35,16 +35,21 @@
         @test node_features(g123)[:,11:14] ≈ node_features(g2) 
 
         # scalar graph features
-        g1 = GNNGraph(random_regular_graph(10,2), gdata=rand())
-        g2 = GNNGraph(random_regular_graph(4,2), gdata=rand())
-        g3 = GNNGraph(random_regular_graph(4,2), gdata=rand())
+        g1 = GNNGraph(g1, gdata=rand())
+        g2 = GNNGraph(g2, gdata=rand())
+        g3 = GNNGraph(g3, gdata=rand())
         g123 = Flux.batch([g1, g2, g3])
         @test g123.gdata.u == [g1.gdata.u, g2.gdata.u, g3.gdata.u]
+
+        # Batch of batches
+        g123123 = Flux.batch([g123, g123])
+        @test g123123.graph_indicator == [fill(1, 10); fill(2, 4); fill(3, 7); fill(4, 10); fill(5, 4); fill(6, 7)]
+        @test g123123.num_graphs == 6
     end
 
     @testset "unbatch" begin
-        g1 = rand_graph(10, 20)
-        g2 = rand_graph(5, 10)
+        g1 = rand_graph(10, 20, graph_type=GRAPH_T)
+        g2 = rand_graph(5, 10, graph_type=GRAPH_T)
         g12 = Flux.batch([g1, g2])
         gs = Flux.unbatch([g1,g2])
         @test length(gs) == 2