faster cat_features

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,27 +358,31 @@ 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
 ```
 """
-function Flux.batch(gs::Vector{<:GNNGraph})
-    nodes = [g.num_nodes for g in gs]
-
-    if all(y -> isa(y, COO_T), [g.graph for g in gs] )
-        edge_indices = [edge_index(g) for g in gs]
-        nodesum = cumsum([0, 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 = reduce(vcat, [get_edge_weight(g) for g in gs])
-        w = w isa Vector{Nothing} ? nothing : w 
-        graph = (s, t, w)
-        graph_indicator = vcat([ones_like(ei[1],Int,nodes[ii]) .+ (ii - 1) for (ii,ei) in enumerate(edge_indices)]...)
-    elseif all(y -> isa(y, ADJMAT_T), [g.graph for g in gs] )
-        graph = blockdiag([g.graph for g in gs]...)
-        graph_indicator = vcat([ones_like(graph,Int,nodes[ii]) .+ (ii - 1) for ii in 1:length(nodes)]...)
+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(nodes), 
+        sum(v_num_nodes), 
         sum([g.num_edges for g in gs]),
-        sum([g.num_graphs 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]),

src/GNNGraphs/utils.jl

@@ -27,9 +27,16 @@ function cat_features(x1::NamedTuple, x2::NamedTuple)
     NamedTuple(k => cat_features(getfield(x1,k), getfield(x2,k)) for k in keys(x1))
 end
 
-function cat_features(xs::Vector{<:NamedTuple})
+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"
+    all(y -> y==symbols[1], symbols) || @error "cannot concatenate feature data with different keys"
     length(xs) == 1 && return xs[1] 
 
     # concatenate 
@@ -82,9 +89,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,10 +18,9 @@
     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)
@@ -36,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