From 0b9fbfe0b113ddbb8b27fdc9e3ee70d70edab3a8 Mon Sep 17 00:00:00 2001 From: Casper Date: Tue, 19 Apr 2022 14:27:47 +0200 Subject: [PATCH 1/2] Explicit unbatching --- src/GNNGraphs/transform.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/GNNGraphs/transform.jl b/src/GNNGraphs/transform.jl index 59f58317e..d5e892e90 100644 --- a/src/GNNGraphs/transform.jl +++ b/src/GNNGraphs/transform.jl @@ -433,6 +433,9 @@ function Flux.unbatch(g::GNNGraph) [getgraph(g, i) for i in 1:g.num_graphs] end +function Flux.unbatch(g::GNNGraph, x) + return [x[g.graph_indicator .∈ Ref(i)] for i in 1:g.num_graphs] +end """ getgraph(g::GNNGraph, i; nmap=false) From 61a965796968b2c32277d80c551f87ffdb3d56c4 Mon Sep 17 00:00:00 2001 From: Casper Date: Tue, 24 May 2022 15:27:55 +0200 Subject: [PATCH 2/2] Speedup unbatching --- src/GNNGraphs/transform.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/GNNGraphs/transform.jl b/src/GNNGraphs/transform.jl index d5e892e90..73d23ec60 100644 --- a/src/GNNGraphs/transform.jl +++ b/src/GNNGraphs/transform.jl @@ -433,8 +433,10 @@ function Flux.unbatch(g::GNNGraph) [getgraph(g, i) for i in 1:g.num_graphs] end -function Flux.unbatch(g::GNNGraph, x) - return [x[g.graph_indicator .∈ Ref(i)] for i in 1:g.num_graphs] +function Flux.unbatch(g::GNNGraph, x) + changes = g.graph_indicator[1:end-1] .!= g.graph_indicator[2:end] + indexes = [0; Array(findall(changes)); length(g.graph_indicator)] + return [x[indexes[i]+1:indexes[i+1]] for i in 1:g.num_graphs] end """