diff --git a/CHANGELOG.md b/CHANGELOG.md
index ffdbc7993..1a57d24b9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and the packages adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
Entries link to the pull request that introduced them.
+## GNNGraphs.jl — Unreleased (towards 1.7.0)
+
+**Added**
+- Graphs can be compiled with Reactant.jl: a `GNNGraph` can now be moved to a Reactant device (e.g. `g |> reactant_device()`) and traced. The COO index vectors and the `graph_indicator` of a `GNNGraph` are no longer required to have an `Integer` eltype, so that traced arrays (eltype `TracedRNumber{Int}`) can be stored in them, and a new `Reactant` package extension makes `degree` work inside a trace ([#730]).
+- Added `pad_graph`, which pads a graph to given numbers of nodes, edges and graphs, putting the padding in extra graphs. Padding every mini-batch to the same size lets Reactant.jl compile a model or a training step once and reuse it for all batches ([#730]).
+- Added the `graph_sizes` field to `GNNGraph`, holding the number of nodes in each graph of a batch (`nothing` for a single graph). Its length is the number of graphs, so a model compiled with Reactant.jl now raises an error when called on a batch with a different number of graphs, instead of silently returning wrong results ([#730]).
+
+## GNNlib.jl — Unreleased (towards 1.5.0)
+
+**Added**
+- Added a `Reactant` package extension, so that layers relying on `propagate` (e.g. `GCNConv`) can be compiled with Reactant.jl and differentiated with Enzyme inside the compiled program. As on CUDA, the sparse adjacency-matrix fast paths of `propagate` are bypassed in favour of gather/scatter message passing. Requires GNNGraphs ≥ 1.7 ([#730]).
+
+**Fixed**
+- `reduce_nodes`, `reduce_edges`, `softmax_nodes`, `softmax_edges` and `softmax_edge_neighbors` now scatter into an output of explicit size (the number of graphs, or of nodes) instead of inferring it from the largest index. This makes `GlobalPool` and the attention layers compile with Reactant, and returns the correct number of columns for batched graphs whose trailing graphs have no nodes ([#730]).
+
## GNNGraphs.jl 1.6.0 - 2026-11-04
**Added**
@@ -247,4 +262,5 @@ Lux implementations of the graph convolutional, pooling, and temporal layers
[#706]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/706
[#722]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/722
[#718]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/718
+[#730]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/730
[FluxML/Zygote.jl#1662]: https://github.com/FluxML/Zygote.jl/issues/1662
diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml
index 29c2d2272..7ec5ad83e 100644
--- a/GNNGraphs/Project.toml
+++ b/GNNGraphs/Project.toml
@@ -1,6 +1,6 @@
name = "GNNGraphs"
uuid = "aed8fd31-079b-4b5a-b342-a13352159b8c"
-version = "1.6.0"
+version = "1.7.0-DEV"
authors = ["Carlo Lucibello and contributors"]
[workspace]
@@ -25,12 +25,14 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
[extensions]
GNNGraphsCUDAExt = "CUDA"
GNNGraphsEnzymeCoreExt = "EnzymeCore"
GNNGraphsMooncakeExt = "Mooncake"
+GNNGraphsReactantExt = "Reactant"
GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
[compat]
@@ -47,6 +49,7 @@ Mooncake = "0.5.24"
NNlib = "0.9"
NearestNeighbors = "0.4"
Random = "1"
+Reactant = "0.2"
SimpleWeightedGraphs = "1.4.0"
SparseArrays = "1"
Statistics = "1"
diff --git a/GNNGraphs/ext/GNNGraphsReactantExt.jl b/GNNGraphs/ext/GNNGraphsReactantExt.jl
new file mode 100644
index 000000000..2dd092a1e
--- /dev/null
+++ b/GNNGraphs/ext/GNNGraphsReactantExt.jl
@@ -0,0 +1,19 @@
+module GNNGraphsReactantExt
+
+using GNNGraphs
+using Reactant: TracedRNumber
+
+# Layers compute degrees with the feature eltype, which inside a Reactant trace is
+# `TracedRNumber{T}`. Allocating with that eltype from host index vectors would create a
+# `Vector{TracedRNumber{T}}`, so compute with the underlying number type instead.
+function GNNGraphs._degree((s, t)::Tuple, ::Type{TracedRNumber{T}}, dir::Symbol,
+ edge_weight::Nothing, num_nodes::Int) where {T}
+ return GNNGraphs._degree((s, t), T, dir, edge_weight, num_nodes)
+end
+
+function GNNGraphs._degree((s, t)::Tuple, ::Type{TracedRNumber{T}}, dir::Symbol,
+ edge_weight::AbstractVector, num_nodes::Int) where {T}
+ return GNNGraphs._degree((s, t), T, dir, edge_weight, num_nodes)
+end
+
+end # module
diff --git a/GNNGraphs/src/GNNGraphs.jl b/GNNGraphs/src/GNNGraphs.jl
index fcfae6cfd..bbf679ab7 100644
--- a/GNNGraphs/src/GNNGraphs.jl
+++ b/GNNGraphs/src/GNNGraphs.jl
@@ -84,6 +84,7 @@ export add_nodes,
perturb_edges,
remove_nodes,
ppr_diffusion,
+ pad_graph,
# from MLUtils.jl
batch,
unbatch,
diff --git a/GNNGraphs/src/abstracttypes.jl b/GNNGraphs/src/abstracttypes.jl
index 73146160f..55f5991cf 100644
--- a/GNNGraphs/src/abstracttypes.jl
+++ b/GNNGraphs/src/abstracttypes.jl
@@ -1,10 +1,13 @@
-const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V <: Union{Nothing, AbstractVector}}
+# The index vectors of COO_T and AVecI are not restricted to integer eltypes so that
+# traced arrays (e.g. Reactant's `TracedRArray{Int}`, whose eltype is `TracedRNumber{Int}`)
+# can be stored in a graph.
+const COO_T = Tuple{T, T, V} where {T <: AbstractVector, V <: Union{Nothing, AbstractVector}}
const ADJLIST_T = AbstractVector{T} where {T <: AbstractVector{<:Integer}}
const ADJMAT_T = AbstractMatrix
const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T
-const AVecI = AbstractVector{<:Integer}
+const AVecI = AbstractVector
# All concrete graph types should be subtypes of AbstractGNNGraph{T}.
# GNNGraph and GNNHeteroGraph are the two concrete types.
diff --git a/GNNGraphs/src/convert.jl b/GNNGraphs/src/convert.jl
index f5675ffff..754da10d0 100644
--- a/GNNGraphs/src/convert.jl
+++ b/GNNGraphs/src/convert.jl
@@ -125,7 +125,7 @@ end
function _to_coo_graph(g::GNNGraph{<:SPARSE_T})
s, t, v = findnz(g.graph)
return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs,
- g.graph_indicator, g.ndata, g.edata, g.gdata)
+ g.graph_indicator, g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
function _to_coo_graph(g::GNNGraph{<:ADJMAT_T})
@@ -133,7 +133,7 @@ function _to_coo_graph(g::GNNGraph{<:ADJMAT_T})
s, t, _ = _findnz_idx(A)
v = _edge_values(A, s, t)
return GNNGraph((s, t, v), g.num_nodes, g.num_edges, g.num_graphs,
- g.graph_indicator, g.ndata, g.edata, g.gdata)
+ g.graph_indicator, g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
### CONVERT TO ADJACENCY MATRIX ################
diff --git a/GNNGraphs/src/gatherscatter.jl b/GNNGraphs/src/gatherscatter.jl
index e897399ed..05aff8295 100644
--- a/GNNGraphs/src/gatherscatter.jl
+++ b/GNNGraphs/src/gatherscatter.jl
@@ -11,7 +11,7 @@ _scatter(aggr, src::Dict, idx, n) = Dict([k => _scatter(aggr, v, idx, n) for (k,
function _scatter(aggr,
src::AbstractArray,
- idx::AbstractVector{<:Integer},
+ idx::AbstractVector,
n::Integer)
dstsize = (size(src)[1:(end - 1)]..., n)
return NNlib.scatter(aggr, src, idx; dstsize)
diff --git a/GNNGraphs/src/gnngraph.jl b/GNNGraphs/src/gnngraph.jl
index dccaf8d67..c98e26cc8 100644
--- a/GNNGraphs/src/gnngraph.jl
+++ b/GNNGraphs/src/gnngraph.jl
@@ -29,7 +29,8 @@ as well, unless explicitly set by the keyword arguments
A `GNNGraph` can also represent multiple graphs batched together
(see [`MLUtils.batch`](@ref) or [`SparseArrays.blockdiag`](@ref)).
The field `g.graph_indicator` contains the graph membership
-of each node.
+of each node, and `g.graph_sizes` the number of nodes in each graph.
+Both are `nothing` for a single graph.
`GNNGraph`s are always directed graphs, therefore each edge is defined
by a source node and a target node (see [`edge_index`](@ref)).
@@ -110,13 +111,16 @@ struct GNNGraph{T <: Union{COO_T, ADJMAT_T}} <: AbstractGNNGraph{T}
num_edges::Int
num_graphs::Int
graph_indicator::Union{Nothing, AVecI} # vector of ints or nothing
+ graph_sizes::Union{Nothing, AVecI} # number of nodes in each graph, or nothing
ndata::DataStore
edata::DataStore
gdata::DataStore
is_coalesced::Bool # only for :coo, true if the graph is coalesced, i.e., no multi edges and indices ordered by target, then source
end
-# GNNGraph constructor setting the is_coalesced field to false
+# GNNGraph constructor deriving the graph_sizes field from the graph indicator.
+# graph_sizes duplicates the graph indicator, but its length is the number of graphs,
+# which makes the number of graphs part of the array shapes seen by compilers like Reactant.
function GNNGraph(graph::T,
num_nodes::Int,
num_edges::Int,
@@ -124,9 +128,16 @@ function GNNGraph(graph::T,
graph_indicator::Union{Nothing, AVecI},
ndata::DataStore,
edata::DataStore,
- gdata::DataStore) where {T <: Union{COO_T, ADJMAT_T}}
+ gdata::DataStore,
+ is_coalesced::Bool = false;
+ graph_sizes = _graph_sizes(graph_indicator, num_graphs)) where {T <: Union{COO_T, ADJMAT_T}}
return GNNGraph{T}(graph, num_nodes, num_edges, num_graphs,
- graph_indicator, ndata, edata, gdata, false)
+ graph_indicator, graph_sizes, ndata, edata, gdata, is_coalesced)
+end
+
+_graph_sizes(graph_indicator::Nothing, num_graphs) = nothing
+function _graph_sizes(graph_indicator::AVecI, num_graphs)
+ return _scatter(+, ones_like(graph_indicator), graph_indicator, num_graphs)
end
function GNNGraph(data::D;
@@ -233,7 +244,7 @@ function GNNGraph(g::GNNGraph; ndata = g.ndata, edata = g.edata, gdata = g.gdata
return GNNGraph(graph,
g.num_nodes, g.num_edges, g.num_graphs,
g.graph_indicator,
- ndata, edata, gdata)
+ ndata, edata, gdata; graph_sizes = g.graph_sizes)
end
"""
@@ -247,12 +258,12 @@ function Base.copy(g::GNNGraph; deep = false)
GNNGraph(deepcopy(g.graph),
g.num_nodes, g.num_edges, g.num_graphs,
deepcopy(g.graph_indicator),
- deepcopy(g.ndata), deepcopy(g.edata), deepcopy(g.gdata))
+ deepcopy(g.ndata), deepcopy(g.edata), deepcopy(g.gdata); graph_sizes = deepcopy(g.graph_sizes))
else
GNNGraph(g.graph,
g.num_nodes, g.num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
end
@@ -345,14 +356,14 @@ MLUtils.getobs(g::GNNGraph, i) = getgraph(g, i)
function Base.:(==)(g1::GNNGraph, g2::GNNGraph)
g1 === g2 && return true
for k in fieldnames(typeof(g1))
- k === :graph_indicator && continue
+ k in (:graph_indicator, :graph_sizes) && continue
getfield(g1, k) != getfield(g2, k) && return false
end
return true
end
function Base.hash(g::T, h::UInt) where {T <: GNNGraph}
- fs = (getfield(g, k) for k in fieldnames(T) if k !== :graph_indicator)
+ fs = (getfield(g, k) for k in fieldnames(T) if k ∉ (:graph_indicator, :graph_sizes))
return foldl((h, f) -> hash(f, h), fs, init = hash(T, h))
end
diff --git a/GNNGraphs/src/sampling.jl b/GNNGraphs/src/sampling.jl
index 88a7b2fea..9ae3084b5 100644
--- a/GNNGraphs/src/sampling.jl
+++ b/GNNGraphs/src/sampling.jl
@@ -202,7 +202,7 @@ function sample_neighbors(g::GNNGraph{<:COO_T}, nodes, K = -1;
gnew = GNNGraph(graph,
g.num_nodes, num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, edata, g.gdata)
+ g.ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
else
nodes_other = dir == :in ? setdiff(s, nodes) : setdiff(t, nodes)
nodes_all = [nodes; nodes_other]
diff --git a/GNNGraphs/src/transform.jl b/GNNGraphs/src/transform.jl
index a89035de7..66fecf0fd 100644
--- a/GNNGraphs/src/transform.jl
+++ b/GNNGraphs/src/transform.jl
@@ -24,7 +24,7 @@ function add_self_loops(g::GNNGraph{<:COO_T})
return GNNGraph((s, t, ew),
g.num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
function add_self_loops(g::GNNGraph{<:ADJMAT_T})
@@ -35,7 +35,7 @@ function add_self_loops(g::GNNGraph{<:ADJMAT_T})
return GNNGraph(A,
g.num_nodes, num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
"""
@@ -60,7 +60,7 @@ function remove_self_loops(g::GNNGraph{<:COO_T})
GNNGraph((s, t, w),
g.num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- g.ndata, edata, g.gdata)
+ g.ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
end
function remove_self_loops(g::GNNGraph{<:ADJMAT_T})
@@ -74,7 +74,7 @@ function remove_self_loops(g::GNNGraph{<:ADJMAT_T})
return GNNGraph(A,
g.num_nodes, num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
"""
@@ -137,7 +137,7 @@ function remove_edges(g::GNNGraph{<:COO_T}, edges_to_remove::AbstractVector{<:In
return GNNGraph((s, t, w),
g.num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- g.ndata, edata, g.gdata)
+ g.ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
end
@@ -185,7 +185,7 @@ function Base.coalesce(g::GNNGraph{<:COO_T}; aggr = +)
return GNNGraph((s, t, w),
g.num_nodes, num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, edata, g.gdata, true)
+ g.ndata, edata, g.gdata, true; graph_sizes = g.graph_sizes)
end
"""
@@ -252,7 +252,7 @@ function remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove::AbstractVector)
return GNNGraph((s, t, w),
num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- ndata, edata, g.gdata)
+ ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
end
"""
@@ -353,7 +353,7 @@ function add_edges(g::GNNGraph{<:COO_T}, data::COO_T; edata = nothing)
return GNNGraph((s, t, w),
num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- ndata, edata, g.gdata)
+ ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
end
"""
@@ -507,7 +507,7 @@ function to_bidirected(g::GNNGraph{<:COO_T})
g = GNNGraph((snew, tnew, w),
g.num_nodes, length(snew), g.num_graphs,
g.graph_indicator,
- g.ndata, edata, g.gdata)
+ g.ndata, edata, g.gdata; graph_sizes = g.graph_sizes)
return coalesce(g; aggr = mean)
end
@@ -527,7 +527,7 @@ function to_unidirected(g::GNNGraph{<:COO_T})
g = GNNGraph((snew, tnew, w),
g.num_nodes, g.num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
return coalesce(g; aggr = mean)
end
@@ -561,7 +561,92 @@ function add_nodes(g::GNNGraph{<:COO_T}, n::Integer; ndata = (;))
GNNGraph(g.graph,
g.num_nodes + n, g.num_edges, g.num_graphs,
g.graph_indicator,
- ndata, g.edata, g.gdata)
+ ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
+end
+
+"""
+ pad_graph(g::GNNGraph, num_nodes, num_edges, num_graphs = g.num_graphs + 1)
+
+Pad `g` to exactly `num_nodes` nodes, `num_edges` edges and `num_graphs` graphs.
+
+All padding nodes and edges are put in one extra graph, appended after the graphs of `g`.
+The padding edges are self-loops on the first padding node. Empty graphs are then
+appended until there are `num_graphs` graphs. Node, edge and graph features, and edge
+weights, are padded with zeros.
+
+Padding every mini-batch to the same size lets a compiler like Reactant.jl compile a
+model, or a training step, once and reuse it for all batches. The original graphs are
+the first `g.num_graphs` ones: mask out the others, e.g. in the loss. A compiled function
+raises an error when called on a graph padded to different sizes, including a different
+number of graphs.
+
+Requires `num_graphs > g.num_graphs`, and `num_nodes > g.num_nodes` if edges are added.
+
+# Examples
+
+```jldoctest
+julia> g = MLUtils.batch([rand_graph(3, 4), rand_graph(5, 6)])
+GNNGraph:
+ num_nodes: 8
+ num_edges: 10
+ num_graphs: 2
+
+julia> gp = pad_graph(g, 10, 12, 4)
+GNNGraph:
+ num_nodes: 10
+ num_edges: 12
+ num_graphs: 4
+
+julia> graph_indicator(gp)
+10-element Vector{Int64}:
+ 1
+ 1
+ 1
+ 2
+ 2
+ 2
+ 2
+ 2
+ 3
+ 3
+```
+"""
+function pad_graph(g::GNNGraph{<:COO_T}, num_nodes::Integer, num_edges::Integer,
+ num_graphs::Integer = g.num_graphs + 1)
+ npad = num_nodes - g.num_nodes
+ epad = num_edges - g.num_edges
+ gpad = num_graphs - g.num_graphs
+ npad >= 0 ||
+ throw(ArgumentError("cannot pad a graph with $(g.num_nodes) nodes to $num_nodes nodes"))
+ epad >= 0 ||
+ throw(ArgumentError("cannot pad a graph with $(g.num_edges) edges to $num_edges edges"))
+ gpad >= 1 ||
+ throw(ArgumentError("`num_graphs` must exceed the $(g.num_graphs) graphs of `g`, to make room for the padding graph"))
+ (epad == 0 || npad >= 1) ||
+ throw(ArgumentError("padding edges need at least one padding node: `num_nodes` must exceed $(g.num_nodes)"))
+
+ s, t, w = edge_index(g)..., get_edge_weight(g)
+ s = vcat(s, fill_like(s, g.num_nodes + 1, epad))
+ t = vcat(t, fill_like(t, g.num_nodes + 1, epad))
+ w = isnothing(w) ? nothing : vcat(w, zeros_like(w, epad))
+ gi = graph_indicator(g)
+ gi = vcat(gi, fill_like(gi, g.num_graphs + 1, npad))
+
+ return GNNGraph((s, t, w),
+ num_nodes, num_edges, num_graphs,
+ gi,
+ _pad_features(g.ndata, npad),
+ _pad_features(g.edata, epad),
+ _pad_features(g.gdata, gpad))
+end
+
+function _pad_features(ds::DataStore, n::Integer)
+ n0 = getn(ds) >= 0 ? getn(ds) : 1
+ return DataStore(n0 + n, Dict{Symbol, Any}(k => _pad_features(x, n) for (k, x) in ds))
+end
+
+function _pad_features(x::AbstractArray, n::Integer)
+ return cat(x, fill!(similar(x, size(x)[1:(end - 1)]..., n), 0); dims = ndims(x))
end
"""
@@ -577,7 +662,7 @@ function set_edge_weight(g::GNNGraph, w::AbstractVector)
return GNNGraph((s, t, w),
g.num_nodes, g.num_edges, g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
function SparseArrays.blockdiag(g1::GNNGraph, g2::GNNGraph)
@@ -1076,5 +1161,5 @@ function ppr_diffusion(g::GNNGraph{<:COO_T}; alpha = 0.85f0)
return GNNGraph((s, t, new_w),
g.num_nodes, length(s), g.num_graphs,
g.graph_indicator,
- g.ndata, g.edata, g.gdata)
+ g.ndata, g.edata, g.gdata; graph_sizes = g.graph_sizes)
end
diff --git a/GNNGraphs/test/gnngraph.jl b/GNNGraphs/test/gnngraph.jl
index ca9c8af02..ff4cfa804 100644
--- a/GNNGraphs/test/gnngraph.jl
+++ b/GNNGraphs/test/gnngraph.jl
@@ -413,6 +413,36 @@ end
end
end
+@testitem "graph_sizes" setup=[GraphsTestModule] begin
+ using .GraphsTestModule
+ @test rand_graph(3, 4).graph_sizes === nothing
+ for GRAPH_T in GRAPH_TYPES
+ g = MLUtils.batch([rand_graph(n, m, graph_type = GRAPH_T) for (n, m) in [(3, 4), (5, 6)]])
+ @test g.graph_sizes == [3, 5]
+ @test add_self_loops(g).graph_sizes == [3, 5]
+ @test GNNGraph(g, ndata = rand(2, 8)).graph_sizes == [3, 5]
+ @test copy(g, deep = true).graph_sizes == [3, 5]
+ @test getgraph(g, 2).graph_sizes == [5]
+ end
+ g = MLUtils.batch([rand_graph(3, 4), rand_graph(5, 6)])
+ @test pad_graph(g, 10, 12, 4).graph_sizes == [3, 5, 2, 0]
+ # the nodes of a graph need not be contiguous
+ g = GNNGraph([1, 2, 3], [2, 3, 4], graph_indicator = [1, 2, 1, 2])
+ @test g.graph_sizes == [2, 2]
+end
+
+@testitem "graph_sizes GPU" setup=[GraphsTestModule] tags=[:gpu] begin
+ using .GraphsTestModule
+ dev = gpu_device(force=true)
+ gs = [rand_graph(3, 4), rand_graph(5, 6)]
+ g = MLUtils.batch(dev.(gs))
+ @test get_device(g.graph_sizes) isa AbstractGPUDevice
+ @test Array(g.graph_sizes) == [3, 5]
+ @test Array(dev(MLUtils.batch(gs)).graph_sizes) == [3, 5]
+ @test Array(add_self_loops(g).graph_sizes) == [3, 5]
+ @test Array(pad_graph(g, 10, 12, 4).graph_sizes) == [3, 5, 2, 0]
+end
+
@testitem "show" begin
# no throw when global data is not an array
g = rand_graph(10, 20, gdata = "ciao")
diff --git a/GNNGraphs/test/transform.jl b/GNNGraphs/test/transform.jl
index 508cc6b16..8bb340cad 100644
--- a/GNNGraphs/test/transform.jl
+++ b/GNNGraphs/test/transform.jl
@@ -414,6 +414,61 @@ end
end
end
+@testitem "pad_graph" setup=[GraphsTestModule] begin
+ using .GraphsTestModule
+ gs = [rand_graph(n, m, ndata = rand(Float32, 2, n), edata = rand(Float32, 3, m),
+ gdata = rand(Float32, 4)) for (n, m) in [(3, 4), (5, 6)]]
+ gs = [set_edge_weight(g, rand(Float32, g.num_edges)) for g in gs]
+ g = MLUtils.batch(gs)
+ gp = pad_graph(g, 10, 13, 4)
+
+ @test gp.num_nodes == 10
+ @test gp.num_edges == 13
+ @test gp.num_graphs == 4
+ @test graph_indicator(gp) == [1, 1, 1, 2, 2, 2, 2, 2, 3, 3]
+ s, t = edge_index(gp)
+ @test s[11:13] == t[11:13] == [9, 9, 9]
+ @test get_edge_weight(gp) == [get_edge_weight(g); zeros(Float32, 3)]
+ @test gp.ndata.x == [g.ndata.x zeros(Float32, 2, 2)]
+ @test gp.edata.e == [g.edata.e zeros(Float32, 3, 3)]
+ @test gp.gdata.u == [g.gdata.u zeros(Float32, 4, 2)]
+
+ # the original graphs are unchanged, and the extra ones are the padding
+ gsp = MLUtils.unbatch(gp)
+ @test gsp[1:2] == MLUtils.unbatch(g)
+ @test gsp[3].num_nodes == 2 && gsp[3].num_edges == 3
+ @test gsp[4].num_nodes == 0 && gsp[4].num_edges == 0
+
+ # a single graph, padded with nodes only
+ g1 = rand_graph(4, 6)
+ gp1 = pad_graph(g1, 5, 6)
+ @test gp1.num_graphs == 2
+ @test graph_indicator(gp1) == [1, 1, 1, 1, 2]
+ @test edge_index(gp1) == edge_index(g1)
+
+ @test_throws ArgumentError pad_graph(g, 7, 13, 4) # fewer nodes
+ @test_throws ArgumentError pad_graph(g, 10, 9, 4) # fewer edges
+ @test_throws ArgumentError pad_graph(g, 10, 13, 2) # no room for the padding graph
+ @test_throws ArgumentError pad_graph(g, 8, 13, 4) # padding edges without padding nodes
+end
+
+@testitem "pad_graph GPU" setup=[GraphsTestModule] tags=[:gpu] begin
+ using .GraphsTestModule
+ dev = gpu_device(force=true)
+ g = MLUtils.batch([rand_graph(n, m, ndata = rand(Float32, 2, n), edata = rand(Float32, 3, m))
+ for (n, m) in [(3, 4), (5, 6)]])
+ gp = pad_graph(g, 10, 13, 4)
+ gp_gpu = pad_graph(g |> dev, 10, 13, 4)
+
+ @test get_device(gp_gpu.graph_indicator) isa AbstractGPUDevice
+ @test get_device(edge_index(gp_gpu)[1]) isa AbstractGPUDevice
+ @test Array(gp_gpu.graph_indicator) == gp.graph_indicator
+ @test Array(edge_index(gp_gpu)[1]) == edge_index(gp)[1]
+ @test Array(edge_index(gp_gpu)[2]) == edge_index(gp)[2]
+ @test Array(gp_gpu.ndata.x) == gp.ndata.x
+ @test Array(gp_gpu.edata.e) == gp.edata.e
+end
+
@testitem "remove_self_loops" setup=[GraphsTestModule] begin
using .GraphsTestModule
for GRAPH_T in GRAPH_TYPES
diff --git a/GNNLux/docs/Project.toml b/GNNLux/docs/Project.toml
index 4d0536640..b4183cd57 100644
--- a/GNNLux/docs/Project.toml
+++ b/GNNLux/docs/Project.toml
@@ -3,11 +3,13 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
GNNLux = "e8545f4d-a905-48ac-a8c4-ca114b98986d"
GNNlib = "a6a84749-d869-43f8-aacc-be26a1996e48"
GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
@@ -15,8 +17,10 @@ MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
TSne = "24678dba-d5e9-5843-a4c6-250288b04835"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
[sources]
GNNGraphs = {path = "../../GNNGraphs"}
@@ -27,3 +31,4 @@ GNNlib = {path = "../../GNNlib"}
DocumenterInterLinks = "1.1.0"
MLDatasets = "0.7"
OneHotArrays = "0.2.11"
+Reactant = "0.2.305"
diff --git a/GNNLux/docs/src/tutorials/beginner_tutorials/gnn_intro.md b/GNNLux/docs/src/tutorials/beginner_tutorials/gnn_intro.md
index 6c488bf11..862dbabc5 100644
--- a/GNNLux/docs/src/tutorials/beginner_tutorials/gnn_intro.md
+++ b/GNNLux/docs/src/tutorials/beginner_tutorials/gnn_intro.md
@@ -4,7 +4,7 @@ EditURL = "../../../src_tutorials/gnn_intro.jl"
# Hands-on introduction to Graph Neural Networks
-*This Pluto notebook is a Julia adaptation of the Pytorch Geometric tutorials that can be found [here](https://pytorch-geometric.readthedocs.io/en/latest/notes/colabs.html).*
+*This tutorial is a Julia adaptation of one of the [Pytorch Geometric tutorials](https://pytorch-geometric.readthedocs.io/en/latest/notes/colabs.html).*
@@ -28,14 +28,28 @@ using MLDatasets
using LinearAlgebra, Random, Statistics
import GraphMakie
import CairoMakie as Makie
-using Zygote, Optimisers, OneHotArrays
+using Reactant, Enzyme, Optimisers, OneHotArrays
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
rng = Random.seed!(10) # for reproducibility
````
````
-TaskLocalRNG()
+Random.TaskLocalRNG()
+````
+
+Following the standard practice in Lux, we compile our models with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate them with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+Reactant traces the Julia code of the model into an [XLA](https://openxla.org/) program and runs it on the best available device: a GPU if one is functional, the CPU otherwise
+(call `Reactant.set_default_backend("cpu")` before the device is created to force the CPU).
+Data, graphs, parameters and states are moved to that device with `xdev`, and results are brought back to the host with `cdev`:
+
+````julia
+const xdev = reactant_device()
+const cdev = cpu_device()
+````
+
+````
+(::MLDataDevices.CPUDevice{Missing}) (generic function with 1 method)
````
Following [Kipf et al. (2017)](https://arxiv.org/abs/1609.02907), let's dive into the world of GNNs by looking at a simple graph-structured example, the well-known [**Zachary's karate club network**](https://en.wikipedia.org/wiki/Zachary%27s_karate_club). This graph describes a social network of 34 members of a karate club and documents links between members who interacted outside the club. Here, we are interested in detecting communities that arise from the member's interaction.
@@ -123,9 +137,9 @@ GNNGraph:
num_nodes: 34
num_edges: 156
ndata:
- y = 4×34 OneHotMatrix(::Vector{UInt32}) with eltype Bool
- train_mask = 34-element Vector{Bool}
- x = 34×34 Matrix{Float32}
+ y = 4×34 OneHotMatrix(::Vector{UInt32}) with eltype Bool
+ train_mask = 34-element Vector{Bool}
+ x = 34×34 Matrix{Float32}
````
Let's now look at the underlying graph in more detail:
@@ -191,7 +205,7 @@ GraphMakie.graphplot(g |> to_unidirected, node_size = 20, node_color = labels, a
````
```@raw html
-
+
```
## Implementing Graph Neural Networks
@@ -234,15 +248,14 @@ function GCN(num_features, num_classes, hidden_dim1, hidden_dim2; use_bias = tru
end
function (gcn::GCN)(g::GNNGraph, x, ps, st) # forward pass
- dense = StatefulLuxLayer{true}(gcn.dense, ps.dense, GNNLux._getstate(st, :dense))
x, stconv1 = gcn.conv1(g, x, ps.conv1, st.conv1)
x = tanh.(x)
x, stconv2 = gcn.conv2(g, x, ps.conv2, st.conv2)
x = tanh.(x)
x, stconv3 = gcn.conv3(g, x, ps.conv3, st.conv3)
x = tanh.(x)
- out = dense(x)
- return (out, x), (conv1 = stconv1, conv2 = stconv2, conv3 = stconv3)
+ out, stdense = gcn.dense(x, ps.dense, st.dense)
+ return (out, x), (conv1 = stconv1, conv2 = stconv2, conv3 = stconv3, dense = stdense)
end
@@ -278,20 +291,35 @@ hidden_dim2 = 2
gcn = GCN(num_features, num_classes, hidden_dim1, hidden_dim2)
ps, st = LuxCore.setup(rng, gcn)
+````
-(ŷ, emb_init), st = gcn(g, g.x, ps, st)
+````
+((conv1 = (weight = Float32[-0.020002399 0.047499612 -0.27350742 -0.2503829 0.19407885 -0.24507841 -0.01664555 -0.060577236 -0.11866234 0.034963984 -0.20158602 -0.18895462 0.32940522 -0.046229932 -0.18350486 0.23901074 -0.04734784 -0.26652414 0.35169592 -0.2931787 0.3739917 0.35536894 0.288215 -0.39124 -0.16256215 -0.21504663 0.038279288 0.18286245 0.08961711 0.23207654 -0.03708796 -0.039563365 -0.22244443 -0.38602307; 0.2995324 0.14075159 -0.30427432 0.13097543 0.28904444 0.21731475 -0.09985288 0.27107894 -0.007117045 0.15674837 -0.022068018 0.042580485 -0.10613822 0.20189701 -0.30136207 -0.3401017 0.16348779 -0.30543056 -0.12484745 -0.07398246 0.26353332 -0.38557786 -0.3389712 -0.0051687597 -0.033157 0.295797 0.05920496 0.3746917 -0.2202008 -0.17036581 -0.09505441 -0.054049794 0.07730663 0.31493124; -0.13803616 -0.03736104 -0.35226765 0.22023676 0.3941705 0.24844564 -0.020854851 0.086532585 0.21445169 -0.3289619 -0.15910763 -0.050556377 -0.15140268 -0.014706691 0.19785336 0.018018777 0.1937682 -0.21754506 0.28674147 0.054676246 0.3883246 0.22768302 -0.15534402 0.3843201 0.32648435 0.36596268 -0.14410815 0.1996352 -0.0016028711 0.2056321 -0.21145336 -0.2560294 -0.04409667 -0.22138521; 0.17314902 -0.18058774 -0.22752328 -0.24605872 0.17544888 -0.2618836 0.3484168 -0.2800446 0.29118395 0.09338584 0.212137 -0.07234587 0.37416515 0.2621097 -0.16015078 -0.24352941 0.16808778 0.3366308 0.30960536 -0.33132365 -0.00071181345 -0.15665922 0.32227686 -0.12654486 -0.3390546 -0.10491658 -0.26917255 -0.26943693 -0.34751803 0.049964264 -0.01279019 -0.042032473 0.14342631 -0.24292947], bias = Float32[0.0, 0.0, 0.0, 0.0]), conv2 = (weight = Float32[0.1678357 -0.7359103 0.7806238 -0.21075805; -0.1638795 0.21836422 -0.7440842 0.51125574; 0.35992146 0.8652138 0.6191888 -0.17565744; -0.70409524 0.82704705 -0.5325756 0.7991618], bias = Float32[0.0, 0.0, 0.0, 0.0]), conv3 = (weight = Float32[0.085344195 0.5199404 0.1966759 -0.030344963; 0.01674807 0.874982 0.6283916 -0.90707684], bias = Float32[0.0, 0.0]), dense = (weight = Float32[-0.049724817 -0.85875297; 0.054479122 0.29839325; -0.20621848 -0.4934963; -0.9001368 0.05491519], bias = Float32[0.0, 0.0, 0.0, 0.0])), (conv1 = NamedTuple(), conv2 = NamedTuple(), conv3 = NamedTuple(), dense = NamedTuple()))
+````
+
+Before running the model, we move the graph (including its node features), the parameters and the states to the device.
+Then we compile the forward pass with `Reactant.@compile`, and call the compiled function as we would call the model itself.
+The result is brought back to the host for plotting.
+
+````julia
+g_dev = g |> xdev
+ps_dev, st_dev = (ps, st) |> xdev
+
+gcn_compiled = @compile gcn(g_dev, g_dev.x, ps_dev, st_dev)
+(ŷ, emb_init), _ = gcn_compiled(g_dev, g_dev.x, ps_dev, st_dev)
+emb_init = cdev(emb_init)
function visualize_embeddings(h; colors = nothing)
xs = h[1, :] |> vec
ys = h[2, :] |> vec
- Makie.scatter(xs, ys, color = labels, markersize = 20)
+ Makie.scatter(xs, ys, color = colors, markersize = 20)
end
visualize_embeddings(emb_init, colors = labels)
````
```@raw html
-
+
```
Remarkably, even before training the weights of our model, the model produces an embedding of nodes that closely resembles the community-structure of the graph.
@@ -311,67 +339,73 @@ After that, we perform multiple rounds of optimization, where each round consist
If you are not new to Lux, this scheme should appear familiar to you.
Note that our semi-supervised learning scenario is achieved by the following line:
-```
-logitcrossentropy(ŷ[:,train_mask], y[:,train_mask])
+```julia
+logitcrossentropy(ŷ[:, train_idx], y[:, train_idx])
```
While we compute node embeddings for all of our nodes, we **only make use of the training nodes for computing the loss**.
-Here, this is implemented by filtering the output of the classifier `out` and ground-truth labels `data.y` to only contain the nodes in the `train_mask`.
+Here, this is implemented by filtering the output of the classifier `ŷ` and ground-truth labels `y` to only contain the training nodes.
+We select them by their integer indices `train_idx = findall(train_mask)` rather than by the mask itself: Reactant currently gives wrong results when a boolean mask indexes the one-hot labels ([Reactant.jl#3087](https://github.com/EnzymeAD/Reactant.jl/issues/3087)).
Let us now start training and see how our node embeddings evolve over time (best experienced by explicitly running the code):
+We use `Lux.Training.single_train_step!` with the `AutoEnzyme()` backend. Since the parameters live on a Reactant device, the first call compiles the whole training step (forward pass, Enzyme gradient and Adam update) into a single XLA program, which is then reused in all following iterations.
+The trained parameters and states are returned in the `train_state`.
+
````julia
-function custom_loss(gcn, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
- (ŷ, _) ,st = gcn(g, x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+function custom_loss(gcn, ps, st, (g, x, y, train_idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
+ (ŷ, _), st = gcn(g, x, ps, st)
+ return logitcrossentropy(ŷ[:, train_idx], y[:, train_idx]), st, (;)
end
-function train_model!(gcn, ps, st, g)
- train_state = Lux.Training.TrainState(gcn, ps, st, Adam(1e-2))
+function train_model(gcn, ps, st, g, train_idx)
+ train_state = Training.TrainState(gcn, ps, st, Adam(1.0f-2))
for iter in 1:2000
- _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), custom_loss,(g, g.x, g.y), train_state)
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss,
+ (g, g.x, g.y, train_idx), train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss)")
+ println("Epoch: $(iter) Loss: $(Float32(loss))")
end
end
- return gcn, ps, st
+ return train_state.parameters, train_state.states
end
-gcn, ps, st = train_model!(gcn, ps, st, g);
+train_idx = findall(train_mask) |> xdev
+ps_dev, st_dev = train_model(gcn, ps_dev, st_dev, g_dev, train_idx);
````
````
-Epoch: 100 Loss: 0.21202108
+Epoch: 100 Loss: 0.2120211
Epoch: 200 Loss: 0.06440282
Epoch: 300 Loss: 0.03517379
-Epoch: 400 Loss: 0.022871211
+Epoch: 400 Loss: 0.022871241
Epoch: 500 Loss: 0.016305141
Epoch: 600 Loss: 0.012315823
-Epoch: 700 Loss: 0.009680824
+Epoch: 700 Loss: 0.009680854
Epoch: 800 Loss: 0.007835168
-Epoch: 900 Loss: 0.006484421
+Epoch: 900 Loss: 0.006484422
Epoch: 1000 Loss: 0.0054619424
-Epoch: 1100 Loss: 0.004666665
+Epoch: 1100 Loss: 0.0046666646
Epoch: 1200 Loss: 0.004034364
-Epoch: 1300 Loss: 0.0035222156
+Epoch: 1300 Loss: 0.0035222159
Epoch: 1400 Loss: 0.003100942
Epoch: 1500 Loss: 0.0027497704
Epoch: 1600 Loss: 0.002453527
Epoch: 1700 Loss: 0.0022013204
Epoch: 1800 Loss: 0.0019845003
Epoch: 1900 Loss: 0.0017966953
-Epoch: 2000 Loss: 0.0016328939
+Epoch: 2000 Loss: 0.0016328938
````
Train accuracy:
````julia
-(ŷ, emb_final), st = gcn(g, g.x, ps, st)
-mean(onecold(ŷ[:, train_mask]) .== onecold(g.y[:, train_mask]))
+(ŷ, emb_final), _ = gcn_compiled(g_dev, g_dev.x, ps_dev, st_dev)
+ŷ, emb_final = cdev((ŷ, emb_final))
+mean(onecold(ŷ[:, train_mask]) .== onecold(y[:, train_mask]))
````
````
@@ -395,7 +429,7 @@ visualize_embeddings(emb_final, colors = labels)
````
```@raw html
-
+
```
As one can see, our 3-layer GCN model manages to linearly separating the communities and classifying most of the nodes correctly.
diff --git a/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md b/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md
index 9a400b8ed..1701d3eff 100644
--- a/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md
+++ b/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md
@@ -11,6 +11,7 @@ EditURL = "../../../src_tutorials/graph_classification.jl"
In this tutorial session we will have a closer look at how to apply **Graph Neural Networks (GNNs) to the task of graph classification**.
Graph classification refers to the problem of classifying entire graphs (in contrast to nodes), given a **dataset of graphs**, based on some structural graph properties and possibly on some input node features.
Here, we want to embed entire graphs, and we want to embed those graphs in such a way so that they are linearly separable given a task at hand.
+We will use a graph convolutional network to create a vector embedding of the input graph, and the apply a simple linear classification head to perform the final classification.
A common graph classification task is **molecular property prediction**, in which molecules are represented as graphs, and the task may be to infer whether a molecule inhibits HIV virus replication or not.
@@ -21,10 +22,20 @@ Let's import the necessary packages. Then we'll load and inspect one of the smal
using Lux, GNNLux
using MLDatasets, MLUtils
using LinearAlgebra, Random, Statistics
-using Zygote, Optimisers, OneHotArrays
+using Reactant, Enzyme, Optimisers, OneHotArrays
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
-rng = Random.seed!(42); # for reproducibility
+Random.seed!(42); # for reproducibility
+rng = Xoshiro(42); # for the model parameters and states
+````
+
+We compile the model with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate it with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+````julia
+const xdev = reactant_device()
+const cdev = cpu_device();
dataset = TUDataset("MUTAG")
````
@@ -119,7 +130,7 @@ test_loader = DataLoader(test_data, batchsize = 32, shuffle = false)
````
````
-2-element DataLoader(::Tuple{Vector{GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, OneHotArrays.OneHotMatrix{UInt32, Vector{UInt32}}}, batchsize=32)
+2-element DataLoader(data::Tuple{Vector{GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, OneHotArrays.OneHotMatrix{UInt32, Vector{UInt32}}}, batchsize=32)
with first element:
(32-element Vector{GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}}, 2×32 OneHotMatrix(::Vector{UInt32}) with eltype Bool,)
````
@@ -148,7 +159,7 @@ vec_gs, _ = first(train_loader)
````
````
-(GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(11, 22) with x: 7×11 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(19, 44) with x: 7×19 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(16, 36) with x: 7×16 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(19, 42) with x: 7×19 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(13, 26) with x: 7×13 data, GNNGraph(19, 40) with x: 7×19 data, GNNGraph(25, 56) with x: 7×25 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(28, 66) with x: 7×28 data, GNNGraph(19, 40) with x: 7×19 data, GNNGraph(19, 44) with x: 7×19 data, GNNGraph(17, 36) with x: 7×17 data, GNNGraph(12, 24) with x: 7×12 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(27, 66) with x: 7×27 data, GNNGraph(17, 38) with x: 7×17 data, GNNGraph(19, 42) with x: 7×19 data, GNNGraph(17, 36) with x: 7×17 data, GNNGraph(12, 26) with x: 7×12 data, GNNGraph(24, 50) with x: 7×24 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(19, 42) with x: 7×19 data, GNNGraph(11, 22) with x: 7×11 data, GNNGraph(16, 34) with x: 7×16 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(16, 36) with x: 7×16 data], Bool[1 0 0 0 0 0 0 0 1 1 0 0 0 0 0 1 0 1 0 0 0 0 0 0 0 0 1 1 1 1 1 0; 0 1 1 1 1 1 1 1 0 0 1 1 1 1 1 0 1 0 1 1 1 1 1 1 1 1 0 0 0 0 0 1])
+(GNNGraphs.GNNGraph{Tuple{Vector{Int64}, Vector{Int64}, Nothing}}[GNNGraph(16, 34) with x: 7×16 data, GNNGraph(11, 22) with x: 7×11 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(18, 40) with x: 7×18 data, GNNGraph(17, 38) with x: 7×17 data, GNNGraph(18, 38) with x: 7×18 data, GNNGraph(14, 30) with x: 7×14 data, GNNGraph(16, 36) with x: 7×16 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(19, 40) with x: 7×19 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(12, 24) with x: 7×12 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(10, 20) with x: 7×10 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(17, 38) with x: 7×17 data, GNNGraph(13, 28) with x: 7×13 data, GNNGraph(11, 22) with x: 7×11 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(20, 46) with x: 7×20 data, GNNGraph(22, 50) with x: 7×22 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(21, 44) with x: 7×21 data, GNNGraph(14, 28) with x: 7×14 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(24, 50) with x: 7×24 data, GNNGraph(20, 44) with x: 7×20 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(19, 44) with x: 7×19 data, GNNGraph(23, 54) with x: 7×23 data, GNNGraph(19, 44) with x: 7×19 data, GNNGraph(23, 54) with x: 7×23 data], Bool[1 1 0 0 0 0 1 0 0 1 0 0 0 1 0 0 1 1 0 0 0 0 0 1 0 0 0 0 0 0 0 0; 0 0 1 1 1 1 0 1 1 0 1 1 1 0 1 1 0 0 1 1 1 1 1 0 1 1 1 1 1 1 1 1])
````
````julia
@@ -157,11 +168,11 @@ MLUtils.batch(vec_gs)
````
GNNGraph:
- num_nodes: 570
- num_edges: 1254
+ num_nodes: 591
+ num_edges: 1316
num_graphs: 32
ndata:
- x = 7×570 Matrix{Float32}
+ x = 7×591 Matrix{Float32}
````
Each batched graph object is equipped with a **`graph_indicator` vector**, which maps each node to its respective graph in the batch:
@@ -203,85 +214,106 @@ nh = 64
nout = 2
model = create_model(nin, nh, nout)
-ps, st = LuxCore.initialparameters(rng, model), LuxCore.initialstates(rng, model);
+ps, st = Lux.setup(rng, model) |> xdev;
````
-````
-┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
-└ @ LuxCore ~/.julia/packages/LuxCore/GlbG3/src/LuxCore.jl:18
+Here, we again make use of the `GCNConv` with $\mathrm{ReLU}(x) = \max(x, 0)$ activation for obtaining localized node embeddings, before we apply our final classifier on top of a graph readout layer.
+### Padding the mini-batches
+
+Reactant compiles a program for fixed array shapes, and `Lux.Training.single_train_step!` compiles the training step once, on its first call, and reuses it for all subsequent calls.
+A mini-batch of graphs, however, has a different number of nodes and edges every time.
+We therefore pad every mini-batch with [`pad_graph`](@ref) to the same number of nodes, edges and graphs: the padding nodes and edges are put in an extra graph, and empty graphs are added until there are `batchsize + 1` graphs.
+Our node budget fits a batch of the largest graphs, plus one padding node:
+
+````julia
+batchsize = train_loader.batchsize
+max_nodes = batchsize * maximum(g.num_nodes for g in graphs) + 1
+max_edges = batchsize * maximum(g.num_edges for g in graphs)
+max_graphs = batchsize + 1
+
+function pad_batch(gs, y)
+ g = MLUtils.batch(gs)
+ n = g.num_graphs
+ g = pad_graph(g, max_nodes, max_edges, max_graphs)
+ y = [Float32.(y) zeros(Float32, size(y, 1), max_graphs - n)]
+ mask = Float32.((1:max_graphs) .<= n)' # the padding graphs are excluded from the loss
+ return g, y, mask, n
+end;
````
-Here, we again make use of the `GCNConv` with $\mathrm{ReLU}(x) = \max(x, 0)$ activation for obtaining localized node embeddings, before we apply our final classifier on top of a graph readout layer.
+The labels are padded too, and the `mask` restricts the loss to the original graphs.
+Note that the mask is passed to the training step as data, i.e. moved to the device with the other arrays: values that are not device arrays would be fixed when the training step is compiled.
Let's train our network for a few epochs to see how well it performs on the training as well as test set:
````julia
-function custom_loss(model, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+function custom_loss(model, ps, st, (g, x, y, mask))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true), agg = l -> sum(l .* mask) / sum(mask))
st = Lux.trainmode(st)
- ŷ, st = model(g, x, ps, st)
- return logitcrossentropy(ŷ, y), st, 0
+ ŷ, st = model(g, x, ps, st)
+ return logitcrossentropy(ŷ, y), st, (;)
end
-function eval_loss_accuracy(model, ps, st, data_loader)
+function eval_loss_accuracy(model_compiled, ps, st, data_loader)
loss = 0.0
acc = 0.0
ntot = 0
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
for (g, y) in data_loader
- g = MLUtils.batch(g)
- n = length(y)
- ŷ, _ = model(g, g.ndata.x, ps, st)
- loss += logitcrossentropy(ŷ, y) * n
- acc += mean((ŷ .> 0) .== y) * n
+ g, y, _, n = pad_batch(g, y)
+ g = xdev(g)
+ ŷ, _ = model_compiled(g, g.x, ps, st)
+ ŷ, y = cdev(ŷ)[:, 1:n], y[:, 1:n]
+ loss += logitcrossentropy(ŷ, y) * n
+ acc += mean((ŷ .> 0) .== y) * n
ntot += n
end
return (loss = round(loss / ntot, digits = 4),
acc = round(acc * 100 / ntot, digits = 2))
end
-function train_model!(model, ps, st; epochs = 500, infotime = 100)
- train_state = Lux.Training.TrainState(model, ps, st, Adam(1e-2))
+function train_model(model, ps, st; epochs = 500, infotime = 100)
+ g, _, _, _ = pad_batch(first(train_loader)...)
+ g = xdev(g)
+ model_compiled = @compile model(g, g.x, ps, Lux.testmode(st))
- function report(epoch)
- train = eval_loss_accuracy(model, ps, st, train_loader)
+ function report(epoch, ps, st)
st = Lux.testmode(st)
- test = eval_loss_accuracy(model, ps, st, test_loader)
- st = Lux.trainmode(st)
+ train = eval_loss_accuracy(model_compiled, ps, st, train_loader)
+ test = eval_loss_accuracy(model_compiled, ps, st, test_loader)
@info (; epoch, train, test)
end
- report(0)
- for iter in 1:epochs
+
+ train_state = Training.TrainState(model, ps, st, Adam(1.0f-2))
+ report(0, ps, st)
+ for epoch in 1:epochs
for (g, y) in train_loader
- g = MLUtils.batch(g)
- _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), custom_loss,(g, g.ndata.x, y), train_state)
+ g, y, mask, _ = pad_batch(g, y)
+ g, y, mask = (g, y, mask) |> xdev
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss,
+ (g, g.x, y, mask), train_state)
end
- iter % infotime == 0 && report(iter)
+ epoch % infotime == 0 && report(epoch, train_state.parameters, train_state.states)
end
- return model, ps, st
+ return train_state.parameters, train_state.states
end
-model, ps, st = train_model!(model, ps, st);
+ps, st = train_model(model, ps, st);
````
````
-┌ Warning: `training` is set to `Val{true}()` but is not being used within an autodiff call (gradient, jacobian, etc...). This will be slow. If you are using a `Lux.jl` model, set it to inference (test) mode using `LuxCore.testmode`. Reliance on this behavior is discouraged, and is not guaranteed by Semantic Versioning, and might be removed without a deprecation cycle. It is recommended to fix this issue in your code.
-└ @ LuxLib.Utils ~/.julia/packages/LuxLib/ru5RQ/src/utils.jl:314
-┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
-└ @ LuxCore ~/.julia/packages/LuxCore/GlbG3/src/LuxCore.jl:18
-[ Info: (epoch = 0, train = (loss = 0.6934, acc = 51.67), test = (loss = 0.6902, acc = 50.0))
-[ Info: (epoch = 100, train = (loss = 0.3979, acc = 81.33), test = (loss = 0.5769, acc = 69.74))
-[ Info: (epoch = 200, train = (loss = 0.3904, acc = 84.0), test = (loss = 0.6402, acc = 65.79))
-[ Info: (epoch = 300, train = (loss = 0.3813, acc = 85.33), test = (loss = 0.6331, acc = 69.74))
-[ Info: (epoch = 400, train = (loss = 0.3682, acc = 85.0), test = (loss = 0.7273, acc = 69.74))
-[ Info: (epoch = 500, train = (loss = 0.3561, acc = 86.67), test = (loss = 0.6825, acc = 73.68))
+[ Info: (epoch = 0, train = (loss = 0.6708, acc = 69.33), test = (loss = 0.6849, acc = 55.26))
+[ Info: (epoch = 100, train = (loss = 0.4381, acc = 79.67), test = (loss = 0.6617, acc = 73.68))
+[ Info: (epoch = 200, train = (loss = 0.397, acc = 85.0), test = (loss = 0.7596, acc = 68.42))
+[ Info: (epoch = 300, train = (loss = 0.3803, acc = 84.67), test = (loss = 0.7549, acc = 69.74))
+[ Info: (epoch = 400, train = (loss = 0.3644, acc = 86.33), test = (loss = 0.8114, acc = 72.37))
+[ Info: (epoch = 500, train = (loss = 0.3654, acc = 86.0), test = (loss = 0.8797, acc = 71.05))
````
-As one can see, our model reaches around **74% test accuracy**.
+As one can see, our model reaches around **70% test accuracy**.
Reasons for the fluctuations in accuracy can be explained by the rather small dataset (only 38 test graphs), and usually disappear once one applies GNNs to larger datasets.
## (Optional) Exercise
diff --git a/GNNLux/docs/src/tutorials/beginner_tutorials/node_classification.md b/GNNLux/docs/src/tutorials/beginner_tutorials/node_classification.md
index f7bc82ffa..39ae120e7 100644
--- a/GNNLux/docs/src/tutorials/beginner_tutorials/node_classification.md
+++ b/GNNLux/docs/src/tutorials/beginner_tutorials/node_classification.md
@@ -16,11 +16,21 @@ using Lux, GNNLux
using MLDatasets
using Plots, TSne
using Random, Statistics
-using Zygote, Optimisers, OneHotArrays, ConcreteStructs
+using Reactant, Enzyme, Optimisers, OneHotArrays, ConcreteStructs
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
-rng = Random.seed!(17); # for reproducibility
+Random.seed!(17); # for reproducibility
+rng = Xoshiro(17); # for the model parameters and states
+````
+
+We compile the models with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate them with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+````julia
+const xdev = reactant_device()
+const cdev = cpu_device();
````
## Visualize
@@ -59,9 +69,9 @@ dataset.metadata
````
Dict{String, Any} with 3 entries:
+ "num_classes" => 7
"name" => "cora"
"classes" => [1, 2, 3, 4, 5, 6, 7]
- "num_classes" => 7
````
The `graphs` variable contains the graph. The `Cora` dataset contains only 1 graph.
@@ -131,13 +141,7 @@ MLP = Chain(Dense(num_features => hidden_channels, relu),
Dropout(drop_rate),
Dense(hidden_channels => num_classes))
-ps, st = Lux.setup(rng, MLP);
-````
-
-````
-┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
-└ @ LuxCore ~/.julia/packages/LuxCore/GlbG3/src/LuxCore.jl:18
-
+ps, st = Lux.setup(rng, MLP) |> xdev;
````
### Training a Multilayer Perceptron
@@ -148,57 +152,61 @@ Here, we first reduce the 1433-dimensional feature vector to a low-dimensional e
Let's train our simple MLP by following a similar procedure as described in [the first part of this tutorial](https://juliagraphs.org/GraphNeuralNetworks.jl/docs/GNNLux.jl/stable/tutorials/gnn_intro/).
We again make use of the **cross entropy loss** and **Adam optimizer**.
This time, we also define a **`accuracy` function** to evaluate how well our final model performs on the test node set (which labels have not been observed during training).
+The training step is compiled by Reactant the first time `single_train_step!` is called, and the trained parameters and states are returned in the `train_state`.
+The loss selects the training nodes by their integer indices `findall(train_mask)` rather than by the mask itself: Reactant currently gives wrong results when a boolean mask indexes the one-hot labels ([Reactant.jl#3087](https://github.com/EnzymeAD/Reactant.jl/issues/3087)).
+For evaluation, we compile and run the forward pass in one go with `Reactant.@jit`, and compute the accuracy on the host.
````julia
-function loss(model, ps, st, x)
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+function loss(model, ps, st, (x, y, idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
ŷ, st = model(x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+ return logitcrossentropy(ŷ[:, idx], y[:, idx]), st, (;)
end
-function train_model!(MLP, ps, st, x, epochs)
- train_state = Lux.Training.TrainState(MLP, ps, st, Adam(1e-3))
+function train_model(MLP, ps, st, data, epochs)
+ train_state = Training.TrainState(MLP, ps, st, Adam(1.0f-3))
for iter in 1:epochs
- _, loss_value, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss, x, train_state)
+ _, loss_value, _, train_state = Training.single_train_step!(AutoEnzyme(), loss, data, train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss_value)")
+ println("Epoch: $(iter) Loss: $(Float32(loss_value))")
end
end
+ return train_state.parameters, train_state.states
end
function accuracy(model, x, ps, st, y, mask)
st = Lux.testmode(st)
- ŷ, st = model(x, ps, st)
- mean(onecold(ŷ)[mask] .== onecold(y)[mask])
+ ŷ, _ = @jit model(x, ps, st)
+ mean(onecold(cdev(ŷ))[mask] .== onecold(y)[mask])
end
-train_model!(MLP, ps, st, x, 2000)
+x_dev, y_dev, train_idx = (x, y, findall(train_mask)) |> xdev
+
+ps, st = train_model(MLP, ps, st, (x_dev, y_dev, train_idx), 2000);
````
````
-┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
-└ @ LuxCore ~/.julia/packages/LuxCore/GlbG3/src/LuxCore.jl:18
-Epoch: 100 Loss: 0.810594
-Epoch: 200 Loss: 0.48982772
-Epoch: 300 Loss: 0.31716076
-Epoch: 400 Loss: 0.2397098
-Epoch: 500 Loss: 0.20041731
-Epoch: 600 Loss: 0.11589075
-Epoch: 700 Loss: 0.21093586
-Epoch: 800 Loss: 0.18869051
-Epoch: 900 Loss: 0.15322906
-Epoch: 1000 Loss: 0.12451931
-Epoch: 1100 Loss: 0.13396983
-Epoch: 1200 Loss: 0.111468166
-Epoch: 1300 Loss: 0.17113678
-Epoch: 1400 Loss: 0.18155631
-Epoch: 1500 Loss: 0.17731342
-Epoch: 1600 Loss: 0.11386197
-Epoch: 1700 Loss: 0.09408201
-Epoch: 1800 Loss: 0.15806198
-Epoch: 1900 Loss: 0.104388796
-Epoch: 2000 Loss: 0.18465123
+Epoch: 100 Loss: 0.8083341
+Epoch: 200 Loss: 0.4324754
+Epoch: 300 Loss: 0.25322104
+Epoch: 400 Loss: 0.23061328
+Epoch: 500 Loss: 0.15893383
+Epoch: 600 Loss: 0.17034031
+Epoch: 700 Loss: 0.23655905
+Epoch: 800 Loss: 0.19057642
+Epoch: 900 Loss: 0.19333595
+Epoch: 1000 Loss: 0.17188784
+Epoch: 1100 Loss: 0.21955554
+Epoch: 1200 Loss: 0.121264964
+Epoch: 1300 Loss: 0.12474382
+Epoch: 1400 Loss: 0.16333841
+Epoch: 1500 Loss: 0.091047734
+Epoch: 1600 Loss: 0.10720601
+Epoch: 1700 Loss: 0.086015575
+Epoch: 1800 Loss: 0.11173484
+Epoch: 1900 Loss: 0.14874706
+Epoch: 2000 Loss: 0.099525936
````
@@ -206,11 +214,11 @@ After training the model, we can call the `accuracy` function to see how well ou
Here, we are interested in the accuracy of the model, *i.e.*, the ratio of correctly classified nodes:
````julia
-accuracy(MLP, x, ps, st, y, .!train_mask)
+accuracy(MLP, x_dev, ps, st, y, .!train_mask)
````
````
-0.5089563862928349
+0.4945482866043614
````
As one can see, our MLP performs rather bad with only about ~50% test accuracy.
@@ -271,2816 +279,64 @@ Now let's visualize the node embeddings of our **untrained** GCN network.
````julia
gcn = GCN(num_features, num_classes, hidden_channels, drop_rate)
-ps, st = Lux.setup(rng, gcn)
-h_untrained, st = gcn(g, x, ps, st)
-h_untrained = h_untrained |> transpose
+ps, st = Lux.setup(rng, gcn) |> xdev
+g_dev = g |> xdev
+h_untrained, _ = @jit gcn(g_dev, x_dev, ps, Lux.testmode(st))
+h_untrained = h_untrained |> cdev |> transpose
visualize_tsne(h_untrained, g.ndata.targets)
````
```@raw html
-
-
-
+
```
We certainly can do better by training our model.
The training and testing procedure is once again the same, but this time we make use of the node features `x` **and** the graph `g` as input to our GCN model.
````julia
-function loss(gcn, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+function loss(gcn, ps, st, (g, x, y, idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
ŷ, st = gcn(g, x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+ return logitcrossentropy(ŷ[:, idx], y[:, idx]), st, (;)
end
-function train_model!(gcn, ps, st, g, x, y)
- train_state = Lux.Training.TrainState(gcn, ps, st, Adam(1e-2))
+function train_model(gcn, ps, st, data)
+ train_state = Training.TrainState(gcn, ps, st, Adam(1.0f-2))
for iter in 1:2000
- _, loss_value, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss,(g, x, y), train_state)
+ _, loss_value, _, train_state = Training.single_train_step!(AutoEnzyme(), loss, data, train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss_value)")
+ println("Epoch: $(iter) Loss: $(Float32(loss_value))")
end
end
- return gcn, ps, st
+ return train_state.parameters, train_state.states
end
-gcn, ps, st = train_model!(gcn, ps, st, g, x, y);
+ps, st = train_model(gcn, ps, st, (g_dev, x_dev, y_dev, train_idx));
````
````
-┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
-└ @ LuxCore ~/.julia/packages/LuxCore/GlbG3/src/LuxCore.jl:18
-Epoch: 100 Loss: 0.019381031
-Epoch: 200 Loss: 0.017426146
-Epoch: 300 Loss: 0.006051709
-Epoch: 400 Loss: 0.0015434261
-Epoch: 500 Loss: 0.0052008606
-Epoch: 600 Loss: 0.025294377
-Epoch: 700 Loss: 0.0012917791
-Epoch: 800 Loss: 0.005089373
-Epoch: 900 Loss: 0.00912053
-Epoch: 1000 Loss: 0.002442247
-Epoch: 1100 Loss: 0.00024606875
-Epoch: 1200 Loss: 0.00046606906
-Epoch: 1300 Loss: 0.002437515
-Epoch: 1400 Loss: 0.00019191795
-Epoch: 1500 Loss: 0.0056298207
-Epoch: 1600 Loss: 0.00020503976
-Epoch: 1700 Loss: 0.0028860446
-Epoch: 1800 Loss: 0.02319943
-Epoch: 1900 Loss: 0.00030635786
-Epoch: 2000 Loss: 0.00013437525
+Epoch: 100 Loss: 0.018460535
+Epoch: 200 Loss: 0.0069120987
+Epoch: 300 Loss: 0.004057126
+Epoch: 400 Loss: 0.001372528
+Epoch: 500 Loss: 0.005683716
+Epoch: 600 Loss: 0.0063499026
+Epoch: 700 Loss: 0.0011878153
+Epoch: 800 Loss: 0.00022846702
+Epoch: 900 Loss: 0.00080485677
+Epoch: 1000 Loss: 0.0050428165
+Epoch: 1100 Loss: 0.004075307
+Epoch: 1200 Loss: 0.00064730056
+Epoch: 1300 Loss: 9.67041e-5
+Epoch: 1400 Loss: 0.00014688371
+Epoch: 1500 Loss: 0.00013341247
+Epoch: 1600 Loss: 0.00011496399
+Epoch: 1700 Loss: 0.005254709
+Epoch: 1800 Loss: 0.01125948
+Epoch: 1900 Loss: 0.00069569895
+Epoch: 2000 Loss: 7.137828e-5
````
@@ -3089,12 +345,12 @@ Now let's evaluate the loss of our trained GCN.
````julia
function accuracy(model, g, x, ps, st, y, mask)
st = Lux.testmode(st)
- ŷ, st = model(g, x, ps, st)
- mean(onecold(ŷ)[mask] .== onecold(y)[mask])
+ ŷ, _ = @jit model(g, x, ps, st)
+ mean(onecold(cdev(ŷ))[mask] .== onecold(y)[mask])
end
-train_accuracy = accuracy(gcn, g, g.ndata.features, ps, st, y, train_mask)
-test_accuracy = accuracy(gcn, g, g.ndata.features, ps, st, y, .!train_mask)
+train_accuracy = accuracy(gcn, g_dev, x_dev, ps, st, y, train_mask)
+test_accuracy = accuracy(gcn, g_dev, x_dev, ps, st, y, .!train_mask)
println("Train accuracy: $(train_accuracy)")
println("Test accuracy: $(test_accuracy)")
@@ -3102,7 +358,7 @@ println("Test accuracy: $(test_accuracy)")
````
Train accuracy: 1.0
-Test accuracy: 0.7636292834890965
+Test accuracy: 0.7663551401869159
````
@@ -3115,2773 +371,13 @@ We can also verify that once again by looking at the output embeddings of our tr
````julia
st = Lux.testmode(st) # inference mode
-out_trained, st = gcn(g, x, ps, st)
-out_trained = out_trained|> transpose
+out_trained, _ = @jit gcn(g_dev, x_dev, ps, st)
+out_trained = out_trained |> cdev |> transpose
visualize_tsne(out_trained, g.ndata.targets)
````
```@raw html
-
-
-
+
```
## (Optional) Exercises
diff --git a/GNNLux/docs/src/tutorials/intermediate_tutorials/traffic_prediction.md b/GNNLux/docs/src/tutorials/intermediate_tutorials/traffic_prediction.md
index f21345cb2..9d379ebff 100644
--- a/GNNLux/docs/src/tutorials/intermediate_tutorials/traffic_prediction.md
+++ b/GNNLux/docs/src/tutorials/intermediate_tutorials/traffic_prediction.md
@@ -16,7 +16,7 @@ We start by importing the necessary libraries. We use `GNNLux.jl`, `Lux.jl` and
using Lux, GNNLux
using MLDatasets: METRLA
using Statistics, Plots, Random
-using Zygote, Optimisers
+using Reactant, Enzyme, Optimisers
const mae = MAELoss()
@@ -24,6 +24,15 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirm
rng = Random.seed!(42); # for reproducibility
````
+We compile the model with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate it with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+````julia
+const xdev = reactant_device()
+const cdev = cpu_device();
+````
+
## Dataset: METR-LA
We use the `METR-LA` dataset from the paper [Diffusion Convolutional Recurrent Neural Network: Data-driven Traffic Forecasting](https://arxiv.org/pdf/1707.01926.pdf), which contains traffic data from loop detectors in the highway of Los Angeles County. The dataset contains traffic speed data from March 1, 2012 to June 30, 2012. The data is collected every 5 minutes, resulting in 12 observations per hour, from 207 sensors. Each sensor is a node in the graph, and the edge weights are the distances between the sensor locations.
@@ -130,10 +139,10 @@ plot_data(features[1:288],1) # Plot the speed of the first sensor for the first
Now let's construct the static graph, the `train_loader` and `data_loader`.
````julia
-graph = GNNGraph(g.edge_index; edata = g.edge_data, g.num_nodes);
+graph = GNNGraph(g.edge_index; edata = g.edge_data, g.num_nodes) |> xdev;
-train_loader = zip(features[1:288], targets[1:288]); # train on 24 hours
-test_loader = zip(features[289:577], targets[289:577]); # test on next 24 hours
+train_loader = zip(features[1:288], targets[1:288]) |> collect |> xdev; # train on 24 hours
+test_loader = zip(features[289:577], targets[289:577]) |> collect |> xdev; # test on next 24 hours
````
## Model: T-GCN
@@ -182,634 +191,640 @@ Let's look at the output of the model for the first batch of the training data.
In Lux, the parameters and the state of the model are kept separate from the model itself, so we first initialize them.
````julia
-ps, st = Lux.setup(rng, model);
+ps, st = Lux.setup(rng, model) |> xdev;
+````
+
+All the inputs have the same shape, so we compile the forward pass once with `Reactant.@compile` and reuse it for all inputs:
+
+````julia
+model_compiled = @compile model(graph, xdev(features[1]), ps, st)
-ŷ, _ = model(graph, features[1], ps, st);
-ŷ
+ŷ, _ = model_compiled(graph, xdev(features[1]), ps, st);
+cdev(ŷ)
````
````
1×3×207 Array{Float32, 3}:
[:, :, 1] =
- -0.00821716 0.0144252 0.0331109
+ -0.00821717 0.0144252 0.033111
[:, :, 2] =
- -0.00965475 0.0114244 0.0313205
+ -0.00965473 0.0114243 0.0313205
[:, :, 3] =
- -0.021871 -0.00552126 0.00915654
+ -0.021871 -0.00552126 0.00915653
[:, :, 4] =
- -0.0011948 0.0242512 0.0433345
+ -0.00119477 0.0242513 0.0433346
[:, :, 5] =
- -0.00177842 0.0234212 0.0419707
+ -0.0017784 0.0234212 0.0419707
[:, :, 6] =
- -0.00148941 0.0233941 0.0428406
+ -0.00148939 0.0233942 0.0428407
[:, :, 7] =
- 0.000171542 0.0263602 0.0475252
+ 0.000171524 0.0263602 0.0475251
[:, :, 8] =
- -0.0239704 -0.00853027 0.004878
+ -0.0239704 -0.0085302 0.00487801
[:, :, 9] =
- -0.0168092 -0.000113703 0.0207041
+ -0.0168092 -0.000113662 0.0207041
[:, :, 10] =
- -0.00818147 0.0140459 0.0358767
+ -0.00818146 0.0140459 0.0358766
[:, :, 11] =
- -0.00618431 0.0190886 0.0378305
+ -0.00618425 0.0190887 0.0378305
[:, :, 12] =
- -0.0309838 -0.0105868 0.0095427
+ -0.0309838 -0.0105868 0.00954276
[:, :, 13] =
- -0.000436462 0.0265467 0.0472731
+ -0.000436436 0.0265467 0.0472731
[:, :, 14] =
- -0.00863666 0.0150596 0.0368315
+ -0.00863662 0.0150596 0.0368316
[:, :, 15] =
- -0.0117289 0.00825214 0.0348723
+ -0.0117289 0.00825211 0.0348723
[:, :, 16] =
- 0.000949625 0.027768 0.0486651
+ 0.000949666 0.0277681 0.0486652
[:, :, 17] =
- 0.00283761 0.0309462 0.0569695
+ 0.00283761 0.0309461 0.0569696
[:, :, 18] =
- -0.00238523 0.022369 0.0403578
+ -0.00238522 0.0223691 0.0403578
[:, :, 19] =
- -0.0306872 -0.0087631 -0.00276568
+ -0.0306872 -0.00876309 -0.0027657
[:, :, 20] =
- -0.0059651 0.0184265 0.0383082
+ -0.00596509 0.0184266 0.0383083
[:, :, 21] =
- -0.00548585 0.0193472 0.0405481
+ -0.0054858 0.0193472 0.0405482
[:, :, 22] =
- -0.0250331 -0.00580959 0.00109873
+ -0.0250331 -0.00580958 0.00109873
[:, :, 23] =
- -0.00356281 0.0230084 0.0466084
+ -0.00356282 0.0230084 0.0466084
[:, :, 24] =
- -0.00714046 0.0130345 0.0312584
+ -0.00714044 0.0130345 0.0312584
[:, :, 25] =
- -0.0118014 0.0109802 0.0286891
+ -0.0118013 0.0109802 0.0286891
[:, :, 26] =
- -0.0151689 0.00218989 0.0175212
+ -0.0151689 0.00218992 0.0175212
[:, :, 27] =
- -0.00242122 0.0275407 0.0453983
+ -0.00242122 0.0275407 0.0453984
[:, :, 28] =
- -0.0346418 -0.02133 -0.00626384
+ -0.0346417 -0.02133 -0.00626384
[:, :, 29] =
- -0.0146368 0.00559739 0.0229257
+ -0.0146368 0.00559741 0.0229256
[:, :, 30] =
- -0.0164983 0.000262856 0.0112214
+ -0.0164983 0.000262871 0.0112214
[:, :, 31] =
- 0.00149859 0.0272778 0.0527369
+ 0.00149861 0.0272778 0.0527369
[:, :, 32] =
- -0.0110724 0.0105973 0.0257418
+ -0.0110723 0.0105973 0.0257418
[:, :, 33] =
-0.0106025 0.00989559 0.0268476
[:, :, 34] =
- 0.00138379 0.0284617 0.049704
+ 0.00138381 0.0284617 0.049704
[:, :, 35] =
- -0.011029 0.00959434 0.0297334
+ -0.011029 0.00959438 0.0297334
[:, :, 36] =
-0.0277521 -0.00468668 0.00168869
[:, :, 37] =
- -0.0106502 0.0124526 0.0259133
+ -0.0106502 0.0124526 0.0259134
[:, :, 38] =
- -0.0101708 0.0108753 0.0293083
+ -0.0101709 0.0108754 0.0293083
[:, :, 39] =
- -0.00371963 0.0198472 0.0389333
+ -0.00371963 0.0198472 0.0389334
[:, :, 40] =
- -0.0087364 0.013785 0.0318417
+ -0.00873641 0.0137849 0.0318417
[:, :, 41] =
- -0.013489 0.00636882 0.0191159
+ -0.013489 0.0063688 0.019116
[:, :, 42] =
- -0.0033318 0.0221838 0.0422297
+ -0.00333179 0.0221838 0.0422297
[:, :, 43] =
-0.0100745 0.0113125 0.0298342
[:, :, 44] =
- -0.00307915 0.0233894 0.0496784
+ -0.00307906 0.0233894 0.0496785
[:, :, 45] =
- -0.0293883 -0.00140207 0.0153636
+ -0.0293883 -0.00140206 0.0153636
[:, :, 46] =
- -0.0203904 -0.000731792 0.0178434
+ -0.0203904 -0.000731803 0.0178434
[:, :, 47] =
- -0.0242228 -0.0108217 7.44835f-5
+ -0.0242228 -0.0108217 7.45282f-5
[:, :, 48] =
- -0.00794924 0.0155703 0.0320494
+ -0.00794923 0.0155703 0.0320494
[:, :, 49] =
-0.000443846 0.0238108 0.0461237
[:, :, 50] =
- -0.00823136 0.0133004 0.0322236
+ -0.00823137 0.0133004 0.0322237
[:, :, 51] =
- -0.0182157 -0.000923406 0.00874105
+ -0.0182157 -0.00092341 0.0087411
[:, :, 52] =
- -0.0197446 0.0024194 0.0129717
+ -0.0197446 0.00241943 0.0129717
[:, :, 53] =
- -0.00978112 0.0111993 0.026973
+ -0.00978113 0.0111992 0.0269731
[:, :, 54] =
- -0.0142461 0.00751363 0.025807
+ -0.0142461 0.00751366 0.025807
[:, :, 55] =
- -0.0101709 0.0108753 0.0293082
+ -0.0101709 0.0108754 0.0293083
[:, :, 56] =
- -0.0297162 -0.0130692 0.00278344
+ -0.0297162 -0.0130692 0.00278346
[:, :, 57] =
- -0.0109779 0.00445043 0.0219942
+ -0.0109779 0.00445045 0.0219942
[:, :, 58] =
- -0.0100355 0.0117707 0.029253
+ -0.0100355 0.0117707 0.0292529
[:, :, 59] =
- -0.0107693 0.0114266 0.0306436
+ -0.0107694 0.0114266 0.0306435
[:, :, 60] =
- -0.0184939 -0.00159456 0.0146626
+ -0.018494 -0.00159457 0.0146626
[:, :, 61] =
- -0.00412727 0.0208961 0.0406313
+ -0.00412725 0.0208962 0.0406313
[:, :, 62] =
- -0.0133903 0.00800258 0.0177869
+ -0.0133903 0.00800263 0.017787
[:, :, 63] =
- -0.0160819 0.00204028 0.0180206
+ -0.0160819 0.00204031 0.0180205
[:, :, 64] =
- -0.00412727 0.0208961 0.0406313
+ -0.00412725 0.0208962 0.0406313
[:, :, 65] =
- -0.00385667 0.0188402 0.0378552
+ -0.00385666 0.0188402 0.0378552
[:, :, 66] =
- -0.00477769 0.01704 0.0354537
+ -0.00477765 0.0170399 0.0354538
[:, :, 67] =
- -0.02122 -0.00378448 0.00212817
+ -0.02122 -0.0037845 0.00212817
[:, :, 68] =
- -0.0057567 0.0203873 0.0390659
+ -0.00575671 0.0203873 0.039066
[:, :, 69] =
- -0.00440173 0.022102 0.0441296
+ -0.00440171 0.022102 0.0441296
[:, :, 70] =
- -0.0055023 0.0198467 0.0387614
+ -0.00550225 0.0198467 0.0387614
[:, :, 71] =
- -0.00894158 0.014867 0.0314622
+ -0.00894156 0.014867 0.0314622
[:, :, 72] =
-0.00465536 0.0215146 0.0430236
[:, :, 73] =
- -0.00324241 0.0248326 0.0482382
+ -0.00324234 0.0248327 0.0482382
[:, :, 74] =
- -0.00730579 0.0162243 0.0332827
+ -0.00730581 0.0162244 0.0332827
[:, :, 75] =
- -0.0121717 0.00827613 0.0243309
+ -0.0121717 0.00827621 0.024331
[:, :, 76] =
- -0.00820319 0.0151796 0.0315771
+ -0.00820316 0.0151796 0.0315772
[:, :, 77] =
- -0.00511628 0.0193203 0.0399505
+ -0.00511631 0.0193204 0.0399504
[:, :, 78] =
- -0.00664458 0.0168547 0.0434569
+ -0.00664458 0.0168547 0.0434568
[:, :, 79] =
- -0.0212049 -0.00525067 0.00802272
+ -0.021205 -0.00525069 0.00802281
[:, :, 80] =
- -0.0326336 -0.0174345 -0.00179935
+ -0.0326336 -0.0174346 -0.0017993
[:, :, 81] =
- 0.000316545 0.0269748 0.0474231
+ 0.000316545 0.0269748 0.047423
[:, :, 82] =
- -0.00344989 0.022354 0.0433216
+ -0.00344989 0.0223539 0.0433216
[:, :, 83] =
- -0.0018939 0.0250479 0.0508254
+ -0.00189396 0.0250478 0.0508254
[:, :, 84] =
-0.0104273 0.0105077 0.0252024
[:, :, 85] =
- -0.0111967 0.0086285 0.0364158
+ -0.0111967 0.00862854 0.0364157
[:, :, 86] =
- -0.0359868 -0.0197313 -0.00156351
+ -0.0359869 -0.0197312 -0.00156337
[:, :, 87] =
- -0.00905617 0.0120922 0.0267821
+ -0.00905614 0.0120921 0.0267821
[:, :, 88] =
- -0.00705249 0.0131835 0.0311062
+ -0.00705252 0.0131835 0.0311061
[:, :, 89] =
- -0.0102068 0.0108329 0.039408
+ -0.0102068 0.0108329 0.0394081
[:, :, 90] =
-0.00732871 0.0155721 0.0379833
[:, :, 91] =
- -0.00896793 0.0119808 0.0278258
+ -0.00896794 0.0119809 0.0278258
[:, :, 92] =
- 0.000171535 0.0263602 0.0475251
+ 0.000171524 0.0263602 0.0475251
[:, :, 93] =
- -0.0292947 -0.0123728 0.0047844
+ -0.0292947 -0.0123728 0.00478438
[:, :, 94] =
- 0.0018874 0.0290393 0.0512981
+ 0.00188743 0.0290392 0.0512981
[:, :, 95] =
- -0.0095189 0.0127057 0.0287307
+ -0.00951888 0.0127057 0.0287307
[:, :, 96] =
- -0.00920017 0.0153237 0.0363948
+ -0.00920013 0.0153237 0.0363948
[:, :, 97] =
- -0.00675596 0.0171003 0.0365007
+ -0.0067559 0.0171003 0.0365007
[:, :, 98] =
- -0.00516706 0.0195319 0.0380939
+ -0.00516705 0.0195318 0.0380938
[:, :, 99] =
- -0.00468567 0.0212621 0.0432734
+ -0.00468566 0.021262 0.0432734
[:, :, 100] =
- -0.0109695 0.0100322 0.0249101
+ -0.0109695 0.0100323 0.02491
[:, :, 101] =
- -0.00480064 0.0175132 0.0372615
+ -0.00480061 0.0175132 0.0372615
[:, :, 102] =
- -0.00175633 0.021358 0.0421611
+ -0.00175633 0.021358 0.0421612
[:, :, 103] =
- -0.00675705 0.0183506 0.0370809
+ -0.00675706 0.0183507 0.0370808
[:, :, 104] =
- -0.00435386 0.0223792 0.0449273
+ -0.00435384 0.0223792 0.0449274
[:, :, 105] =
- -0.00948677 0.012299 0.0277941
+ -0.00948676 0.012299 0.0277941
[:, :, 106] =
-0.0310867 -0.0244714 -0.0213755
[:, :, 107] =
- -0.039139 -0.0349075 -0.0324736
+ -0.039139 -0.0349075 -0.0324735
[:, :, 108] =
- -0.0197428 -0.00222888 0.0120533
+ -0.0197428 -0.00222884 0.0120533
[:, :, 109] =
- -0.0340344 -0.0180384 -0.00232398
+ -0.0340344 -0.0180385 -0.00232385
[:, :, 110] =
-0.016809 0.0020473 0.0137824
[:, :, 111] =
- -0.0258126 -0.00946338 0.00629207
+ -0.0258126 -0.00946337 0.00629203
[:, :, 112] =
- -0.0113044 0.00965049 0.0271328
+ -0.0113044 0.00965048 0.0271328
[:, :, 113] =
- -0.01398 0.00985853 0.0250579
+ -0.01398 0.00985851 0.0250579
[:, :, 114] =
-0.00451003 0.0224715 0.0425556
[:, :, 115] =
- -0.0107693 0.0114266 0.0306436
+ -0.0107694 0.0114266 0.0306435
[:, :, 116] =
-0.0128593 0.00807549 0.0214224
[:, :, 117] =
- -0.0116444 0.00762197 0.0239993
+ -0.0116444 0.00762193 0.0239992
[:, :, 118] =
-0.0108711 0.0117586 0.0249509
[:, :, 119] =
- -0.00560274 0.0204472 0.0397052
+ -0.00560271 0.0204472 0.0397053
[:, :, 120] =
- -0.0121718 0.00827614 0.0243309
+ -0.0121717 0.00827621 0.024331
[:, :, 121] =
- -0.0151689 0.00218991 0.0175213
+ -0.0151689 0.00218992 0.0175212
[:, :, 122] =
- -0.0158763 0.000667531 0.00323578
+ -0.0158763 0.00066749 0.00323581
[:, :, 123] =
- -0.00835298 0.0145499 0.0314063
+ -0.00835296 0.0145499 0.0314062
[:, :, 124] =
-0.0062982 0.0156706 0.0363288
[:, :, 125] =
- -0.00961787 0.0109065 0.0273044
+ -0.0096179 0.0109065 0.0273045
[:, :, 126] =
- -0.00960421 0.014225 0.036486
+ -0.00960417 0.014225 0.036486
[:, :, 127] =
- -0.0163094 0.00533274 0.00596636
+ -0.0163094 0.00533273 0.0059664
[:, :, 128] =
- -0.00718839 0.0162061 0.0339369
+ -0.00718837 0.0162062 0.0339369
[:, :, 129] =
- -0.00301997 0.0229231 0.0425203
+ -0.00301993 0.0229231 0.0425203
[:, :, 130] =
- -0.00686035 0.0145443 0.0342574
+ -0.00686034 0.0145443 0.0342574
[:, :, 131] =
-0.00408264 0.0218743 0.0489538
[:, :, 132] =
- -0.00594971 0.0171859 0.0346717
+ -0.00594977 0.0171859 0.0346717
[:, :, 133] =
- -0.0358658 -0.0192914 -0.000742409
+ -0.0358658 -0.0192913 -0.00074238
[:, :, 134] =
- -0.00103747 0.0219577 0.0432877
+ -0.00103746 0.0219577 0.0432877
[:, :, 135] =
- -0.0055554 0.016153 0.0345801
+ -0.00555539 0.0161529 0.03458
[:, :, 136] =
- -0.00702389 0.0144618 0.0335251
+ -0.00702387 0.0144618 0.0335252
[:, :, 137] =
- 0.000548251 0.0267444 0.0477673
+ 0.000548206 0.0267443 0.0477674
[:, :, 138] =
- -0.00526601 0.0187382 0.0375164
+ -0.00526599 0.0187381 0.0375164
[:, :, 139] =
- 0.000486601 0.0265969 0.0475849
+ 0.00048662 0.0265969 0.047585
[:, :, 140] =
- -0.0104272 0.0093519 0.0272851
+ -0.0104272 0.00935189 0.0272851
[:, :, 141] =
- -0.00952506 0.0126065 0.0285398
+ -0.00952503 0.0126065 0.0285398
[:, :, 142] =
- -0.00322082 0.0238048 0.0447394
+ -0.00322083 0.0238048 0.0447393
[:, :, 143] =
- -0.0092582 0.0113804 0.0304347
+ -0.0092582 0.0113804 0.0304346
[:, :, 144] =
- -0.0118315 0.00599926 0.0232591
+ -0.0118315 0.00599931 0.0232591
[:, :, 145] =
- 0.000520378 0.0270373 0.0474962
+ 0.000520375 0.0270373 0.0474961
[:, :, 146] =
- -0.0125182 0.00752925 0.0221856
+ -0.0125181 0.00752926 0.0221856
[:, :, 147] =
- -0.00328181 0.0224241 0.0491296
+ -0.00328176 0.0224241 0.0491295
[:, :, 148] =
- -0.0108386 0.008702 0.0250677
+ -0.0108385 0.00870199 0.0250677
[:, :, 149] =
- -0.00445814 0.0179113 0.0386408
+ -0.00445813 0.0179113 0.0386408
[:, :, 150] =
- -0.0143861 0.00381661 0.0164905
+ -0.0143862 0.00381664 0.0164906
[:, :, 151] =
- -0.0108416 0.00402115 0.0214087
+ -0.0108417 0.00402118 0.0214087
[:, :, 152] =
- -0.0102509 0.0109841 0.0389173
+ -0.0102509 0.010984 0.0389172
[:, :, 153] =
- -0.0270753 0.00281014 0.0234369
+ -0.0270753 0.00281012 0.0234368
[:, :, 154] =
- -0.0199297 -0.0041375 0.00616935
+ -0.0199297 -0.00413751 0.00616927
[:, :, 155] =
- -0.00125065 0.0257128 0.0471909
+ -0.00125067 0.0257129 0.0471909
[:, :, 156] =
- -0.00341222 0.0222467 0.0423326
+ -0.0034122 0.0222467 0.0423326
[:, :, 157] =
- -0.00917533 0.0137064 0.0298123
+ -0.00917531 0.0137065 0.0298123
[:, :, 158] =
- -0.00322511 0.0219762 0.0404697
+ -0.00322509 0.0219762 0.0404697
[:, :, 159] =
- -0.006485 0.0166181 0.0342301
+ -0.00648501 0.0166181 0.0342301
[:, :, 160] =
- -0.00353946 0.0222017 0.0419642
+ -0.00353943 0.0222016 0.0419642
[:, :, 161] =
- -0.00112146 0.0259375 0.0471458
+ -0.00112145 0.0259375 0.0471457
[:, :, 162] =
- -0.00145669 0.0243457 0.0434363
+ -0.0014567 0.0243457 0.0434362
[:, :, 163] =
- -0.00114095 0.025566 0.0460631
+ -0.0011409 0.025566 0.0460632
[:, :, 164] =
- -0.00399512 0.0216691 0.0413134
+ -0.00399513 0.0216692 0.0413134
[:, :, 165] =
- -0.00746722 0.0159911 0.0328996
+ -0.0074672 0.0159911 0.0328996
[:, :, 166] =
-0.00272395 0.0207096 0.0416148
[:, :, 167] =
- 0.00169669 0.0287229 0.0532165
+ 0.00169671 0.0287229 0.0532164
[:, :, 168] =
- -0.0178723 0.00496994 0.0235218
+ -0.0178723 0.00496996 0.023522
[:, :, 169] =
-0.0116743 0.0119938 0.0318312
[:, :, 170] =
- -0.011565 0.0115453 0.0313856
+ -0.011565 0.0115452 0.0313855
[:, :, 171] =
- -0.00521046 0.012802 0.0320369
+ -0.00521043 0.012802 0.0320369
[:, :, 172] =
- -0.0187236 0.00331177 0.0242589
+ -0.0187236 0.0033117 0.024259
[:, :, 173] =
- -0.00695409 0.0158522 0.0344385
+ -0.00695409 0.0158522 0.0344386
[:, :, 174] =
- -0.0156356 0.00588394 0.0229958
+ -0.0156356 0.00588392 0.0229959
[:, :, 175] =
- -0.00952948 0.0107777 0.0265258
+ -0.00952949 0.0107777 0.0265258
[:, :, 176] =
- -0.00528941 0.0203201 0.0417827
+ -0.00528939 0.0203201 0.0417828
[:, :, 177] =
- -0.0120879 0.00868366 0.0226738
+ -0.0120879 0.0086836 0.0226737
[:, :, 178] =
- -0.00867432 0.0133468 0.0330589
+ -0.00867435 0.0133468 0.0330589
[:, :, 179] =
- -0.0271077 0.00287797 0.023712
+ -0.0271077 0.00287803 0.023712
[:, :, 180] =
- -0.0292645 -0.0097343 0.00924958
+ -0.0292646 -0.0097343 0.00924959
[:, :, 181] =
- -0.00465781 0.0169174 0.0371693
+ -0.00465783 0.0169174 0.0371692
[:, :, 182] =
- -0.00607961 0.0166836 0.0353746
+ -0.00607962 0.0166835 0.0353746
[:, :, 183] =
- -0.0146344 0.00502735 0.0249425
+ -0.0146344 0.00502737 0.0249425
[:, :, 184] =
- -0.0116239 0.00889217 0.0274019
+ -0.0116239 0.00889223 0.0274018
[:, :, 185] =
- -0.0197955 -0.00350938 0.00836946
+ -0.0197955 -0.00350935 0.00836947
[:, :, 186] =
- -0.0142319 0.00364592 0.0281281
+ -0.0142319 0.00364596 0.028128
[:, :, 187] =
- -0.010957 0.00966091 0.0342825
+ -0.0109571 0.00966085 0.0342825
[:, :, 188] =
-0.00322242 0.0222684 0.0412074
[:, :, 189] =
- 0.00044544 0.0276311 0.0486719
+ 0.000445481 0.0276311 0.0486718
[:, :, 190] =
- -0.0188853 0.00264707 0.00703365
+ -0.0188853 0.00264714 0.00703366
[:, :, 191] =
- -0.00330989 0.024461 0.0483255
+ -0.00330991 0.024461 0.0483256
[:, :, 192] =
- -0.000522427 0.0256213 0.0452741
+ -0.000522409 0.0256213 0.0452741
[:, :, 193] =
-0.00106541 0.0245034 0.0437235
[:, :, 194] =
- -0.00190246 0.0238661 0.0429132
+ -0.0019025 0.0238661 0.0429132
[:, :, 195] =
- -0.0144073 0.00319022 0.0169229
+ -0.0144073 0.00319023 0.0169229
[:, :, 196] =
- -0.000914615 0.0247917 0.0438946
+ -0.000914596 0.0247918 0.0438946
[:, :, 197] =
- 0.00169674 0.0287229 0.0532165
+ 0.00169671 0.0287229 0.0532164
[:, :, 198] =
- -0.000986643 0.0247489 0.0485976
+ -0.000986651 0.0247489 0.0485975
[:, :, 199] =
- -0.00350053 0.0219981 0.0410607
+ -0.0035005 0.0219982 0.0410606
[:, :, 200] =
- -0.00952503 0.0126065 0.0285399
+ -0.00952503 0.0126065 0.0285398
[:, :, 201] =
- -0.0310937 -0.0111156 -0.00414067
+ -0.0310937 -0.0111156 -0.00414072
[:, :, 202] =
- -0.0221025 0.00454363 0.0246241
+ -0.0221025 0.00454361 0.0246241
[:, :, 203] =
-0.0162078 0.00394783 0.0199544
[:, :, 204] =
- -0.0201875 -0.00508371 0.00845758
+ -0.0201875 -0.00508371 0.00845761
[:, :, 205] =
- -0.00533505 0.0184017 0.037045
+ -0.00533503 0.0184017 0.037045
[:, :, 206] =
- -0.0069662 0.0168798 0.0363927
+ -0.00696619 0.0168799 0.0363928
[:, :, 207] =
- -0.00615102 0.0178156 0.0365203
+ -0.00615101 0.0178156 0.0365203
````
The output of the model is a tensor of size `(1, 3, 207)`, which corresponds to the dimension of the feature (in this case speed), the number of time steps, and the number of nodes in the graph, respectively. The model outputs the predicted traffic speed for each sensor at each time step.
@@ -819,6 +834,7 @@ The output of the model is a tensor of size `(1, 3, 207)`, which corresponds to
## Training
We train the model for 100 epochs, using the Adam optimizer with a learning rate of 0.001. We use the mean absolute error (MAE) as the loss function.
+The first call to `single_train_step!` compiles the whole training step (forward pass, Enzyme gradient and Adam update), which is then reused for all the following steps.
````julia
function loss_function(model, ps, st, (graph, x, y))
@@ -828,16 +844,17 @@ end
function train(graph, train_loader, model, ps, st)
- train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001))
+ train_state = Training.TrainState(model, ps, st, Adam(0.001f0))
for epoch in 1:100
+ losses = Float32[]
for (x, y) in train_loader
- _, _, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss_function, (graph, x, y), train_state)
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), loss_function, (graph, x, y), train_state)
+ push!(losses, loss)
end
if epoch % 10 == 0
- ps, st = train_state.parameters, train_state.states
- loss = mean([mae(model(graph, x, ps, st)[1], y) for (x, y) in train_loader])
+ loss = mean(losses)
@show epoch, loss
end
end
@@ -848,16 +865,16 @@ ps, st = train(graph, train_loader, model, ps, st);
````
````
-(epoch, loss) = (10, 0.26200253f0)
-(epoch, loss) = (20, 0.2570168f0)
-(epoch, loss) = (30, 0.253641f0)
-(epoch, loss) = (40, 0.2518906f0)
-(epoch, loss) = (50, 0.25033492f0)
-(epoch, loss) = (60, 0.2518252f0)
-(epoch, loss) = (70, 0.25138408f0)
-(epoch, loss) = (80, 0.24981792f0)
-(epoch, loss) = (90, 0.24679154f0)
-(epoch, loss) = (100, 0.2435512f0)
+(epoch, loss) = (10, 0.25935468f0)
+(epoch, loss) = (20, 0.2550633f0)
+(epoch, loss) = (30, 0.25186265f0)
+(epoch, loss) = (40, 0.24539487f0)
+(epoch, loss) = (50, 0.24060875f0)
+(epoch, loss) = (60, 0.23551488f0)
+(epoch, loss) = (70, 0.22963661f0)
+(epoch, loss) = (80, 0.22238377f0)
+(epoch, loss) = (90, 0.21950722f0)
+(epoch, loss) = (100, 0.21583754f0)
````
@@ -868,7 +885,7 @@ function plot_predicted_data(graph, features, targets, sensor)
ground_truth = []
for i in 1:3:length(features)
push!(ground_truth,targets[i][1,:,sensor])
- push!(prediction, model(graph, features[i], ps, st)[1][1,:,sensor])
+ push!(prediction, cdev(model_compiled(graph, xdev(features[i]), ps, st)[1])[1,:,sensor])
end
prediction = reduce(vcat,prediction)
ground_truth = reduce(vcat, ground_truth)
@@ -881,7 +898,7 @@ plot_predicted_data(graph,features[289:577],targets[289:577], 1)
````
```@raw html
-
+
```
````julia
@@ -895,11 +912,11 @@ accuracy (generic function with 1 method)
Test accuracy:
````julia
-mean([accuracy(model(graph, x, ps, st)[1], y) for (x, y) in test_loader])
+mean([accuracy(cdev(model_compiled(graph, x, ps, st)[1]), cdev(y)) for (x, y) in test_loader])
````
````
-0.6099206f0
+0.6279825f0
````
The accuracy is not very good but can be improved by training using more data. We used a small subset of the dataset for this tutorial because of the computational cost of training the model. From the plot of the predictions, we can see that the model is able to capture the general trend of the traffic speed, but it is not able to capture the peaks of the traffic.
diff --git a/GNNLux/docs/src_tutorials/gnn_intro.jl b/GNNLux/docs/src_tutorials/gnn_intro.jl
index 0b8fbe135..ab397cd5e 100644
--- a/GNNLux/docs/src_tutorials/gnn_intro.jl
+++ b/GNNLux/docs/src_tutorials/gnn_intro.jl
@@ -23,11 +23,19 @@ using MLDatasets
using LinearAlgebra, Random, Statistics
import GraphMakie
import CairoMakie as Makie
-using Zygote, Optimisers, OneHotArrays
+using Reactant, Enzyme, Optimisers, OneHotArrays
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
rng = Random.seed!(10) # for reproducibility
+# Following the standard practice in Lux, we compile our models with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate them with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+# Reactant traces the Julia code of the model into an [XLA](https://openxla.org/) program and runs it on the best available device: a GPU if one is functional, the CPU otherwise
+# (call `Reactant.set_default_backend("cpu")` before the device is created to force the CPU).
+# Data, graphs, parameters and states are moved to that device with `xdev`, and results are brought back to the host with `cdev`:
+
+const xdev = reactant_device()
+const cdev = cpu_device()
+
# Following [Kipf et al. (2017)](https://arxiv.org/abs/1609.02907), let's dive into the world of GNNs by looking at a simple graph-structured example, the well-known [**Zachary's karate club network**](https://en.wikipedia.org/wiki/Zachary%27s_karate_club). This graph describes a social network of 34 members of a karate club and documents links between members who interacted outside the club. Here, we are interested in detecting communities that arise from the member's interaction.
# GNNLux.jl provides utilities to convert [MLDatasets.jl](https://github.com/JuliaML/MLDatasets.jl)'s datasets to its own type:
@@ -149,15 +157,14 @@ function GCN(num_features, num_classes, hidden_dim1, hidden_dim2; use_bias = tru
end
function (gcn::GCN)(g::GNNGraph, x, ps, st) # forward pass
- dense = StatefulLuxLayer{true}(gcn.dense, ps.dense, GNNLux._getstate(st, :dense))
x, stconv1 = gcn.conv1(g, x, ps.conv1, st.conv1)
x = tanh.(x)
x, stconv2 = gcn.conv2(g, x, ps.conv2, st.conv2)
x = tanh.(x)
x, stconv3 = gcn.conv3(g, x, ps.conv3, st.conv3)
- x = tanh.(x)
- out = dense(x)
- return (out, x), (conv1 = stconv1, conv2 = stconv2, conv3 = stconv3)
+ x = tanh.(x)
+ out, stdense = gcn.dense(x, ps.dense, st.dense)
+ return (out, x), (conv1 = stconv1, conv2 = stconv2, conv3 = stconv3, dense = stdense)
end
@@ -194,12 +201,22 @@ hidden_dim2 = 2
gcn = GCN(num_features, num_classes, hidden_dim1, hidden_dim2)
ps, st = LuxCore.setup(rng, gcn)
-(ŷ, emb_init), st = gcn(g, g.x, ps, st)
+
+# Before running the model, we move the graph (including its node features), the parameters and the states to the device.
+# Then we compile the forward pass with `Reactant.@compile`, and call the compiled function as we would call the model itself.
+# The result is brought back to the host for plotting.
+
+g_dev = g |> xdev
+ps_dev, st_dev = (ps, st) |> xdev
+
+gcn_compiled = @compile gcn(g_dev, g_dev.x, ps_dev, st_dev)
+(ŷ, emb_init), _ = gcn_compiled(g_dev, g_dev.x, ps_dev, st_dev)
+emb_init = cdev(emb_init)
function visualize_embeddings(h; colors = nothing)
xs = h[1, :] |> vec
ys = h[2, :] |> vec
- Makie.scatter(xs, ys, color = labels, markersize = 20)
+ Makie.scatter(xs, ys, color = colors, markersize = 20)
end
visualize_embeddings(emb_init, colors = labels)
@@ -223,40 +240,44 @@ visualize_embeddings(emb_init, colors = labels)
# Note that our semi-supervised learning scenario is achieved by the following line:
# ```julia
-# logitcrossentropy(ŷ[:,train_mask], y[:,train_mask])
+# logitcrossentropy(ŷ[:, train_idx], y[:, train_idx])
# ```
# While we compute node embeddings for all of our nodes, we **only make use of the training nodes for computing the loss**.
-# Here, this is implemented by filtering the output of the classifier `out` and ground-truth labels `data.y` to only contain the nodes in the `train_mask`.
+# Here, this is implemented by filtering the output of the classifier `ŷ` and ground-truth labels `y` to only contain the training nodes.
+# We select them by their integer indices `train_idx = findall(train_mask)` rather than by the mask itself: Reactant currently gives wrong results when a boolean mask indexes the one-hot labels ([Reactant.jl#3087](https://github.com/EnzymeAD/Reactant.jl/issues/3087)).
# Let us now start training and see how our node embeddings evolve over time (best experienced by explicitly running the code):
-function custom_loss(gcn, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
- (ŷ, _) ,st = gcn(g, x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+# We use `Lux.Training.single_train_step!` with the `AutoEnzyme()` backend. Since the parameters live on a Reactant device, the first call compiles the whole training step (forward pass, Enzyme gradient and Adam update) into a single XLA program, which is then reused in all following iterations.
+# The trained parameters and states are returned in the `train_state`.
+
+function custom_loss(gcn, ps, st, (g, x, y, train_idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
+ (ŷ, _), st = gcn(g, x, ps, st)
+ return logitcrossentropy(ŷ[:, train_idx], y[:, train_idx]), st, (;)
end
-function train_model!(gcn, ps, st, g)
- train_state = Lux.Training.TrainState(gcn, ps, st, Adam(1e-2))
+function train_model(gcn, ps, st, g, train_idx)
+ train_state = Training.TrainState(gcn, ps, st, Adam(1.0f-2))
for iter in 1:2000
- _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), custom_loss,(g, g.x, g.y), train_state)
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss,
+ (g, g.x, g.y, train_idx), train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss)")
+ println("Epoch: $(iter) Loss: $(Float32(loss))")
end
end
- return gcn, ps, st
+ return train_state.parameters, train_state.states
end
-gcn, ps, st = train_model!(gcn, ps, st, g);
-
-
+train_idx = findall(train_mask) |> xdev
+ps_dev, st_dev = train_model(gcn, ps_dev, st_dev, g_dev, train_idx);
# Train accuracy:
-(ŷ, emb_final), st = gcn(g, g.x, ps, st)
-mean(onecold(ŷ[:, train_mask]) .== onecold(g.y[:, train_mask]))
+(ŷ, emb_final), _ = gcn_compiled(g_dev, g_dev.x, ps_dev, st_dev)
+ŷ, emb_final = cdev((ŷ, emb_final))
+mean(onecold(ŷ[:, train_mask]) .== onecold(y[:, train_mask]))
# Test accuracy:
diff --git a/GNNLux/docs/src_tutorials/graph_classification.jl b/GNNLux/docs/src_tutorials/graph_classification.jl
index c62365736..70785cadd 100644
--- a/GNNLux/docs/src_tutorials/graph_classification.jl
+++ b/GNNLux/docs/src_tutorials/graph_classification.jl
@@ -18,10 +18,18 @@
using Lux, GNNLux
using MLDatasets, MLUtils
using LinearAlgebra, Random, Statistics
-using Zygote, Optimisers, OneHotArrays
+using Reactant, Enzyme, Optimisers, OneHotArrays
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
-rng = Random.seed!(42); # for reproducibility
+Random.seed!(42); # for reproducibility
+rng = Xoshiro(42); # for the model parameters and states
+
+# We compile the model with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate it with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+# Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+# We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+const xdev = reactant_device()
+const cdev = cpu_device();
dataset = TUDataset("MUTAG")
#
@@ -123,62 +131,91 @@ nh = 64
nout = 2
model = create_model(nin, nh, nout)
-ps, st = LuxCore.initialparameters(rng, model), LuxCore.initialstates(rng, model);
+ps, st = Lux.setup(rng, model) |> xdev;
# Here, we again make use of the `GCNConv` with $\mathrm{ReLU}(x) = \max(x, 0)$ activation for obtaining localized node embeddings, before we apply our final classifier on top of a graph readout layer.
+# ### Padding the mini-batches
+
+# Reactant compiles a program for fixed array shapes, and `Lux.Training.single_train_step!` compiles the training step once, on its first call, and reuses it for all subsequent calls.
+# A mini-batch of graphs, however, has a different number of nodes and edges every time.
+# We therefore pad every mini-batch with [`pad_graph`](@ref) to the same number of nodes, edges and graphs: the padding nodes and edges are put in an extra graph, and empty graphs are added until there are `batchsize + 1` graphs.
+# Our node budget fits a batch of the largest graphs, plus one padding node:
+
+batchsize = train_loader.batchsize
+max_nodes = batchsize * maximum(g.num_nodes for g in graphs) + 1
+max_edges = batchsize * maximum(g.num_edges for g in graphs)
+max_graphs = batchsize + 1
+
+function pad_batch(gs, y)
+ g = MLUtils.batch(gs)
+ n = g.num_graphs
+ g = pad_graph(g, max_nodes, max_edges, max_graphs)
+ y = [Float32.(y) zeros(Float32, size(y, 1), max_graphs - n)]
+ mask = Float32.((1:max_graphs) .<= n)' # the padding graphs are excluded from the loss
+ return g, y, mask, n
+end;
+
+# The labels are padded too, and the `mask` restricts the loss to the original graphs.
+# Note that the mask is passed to the training step as data, i.e. moved to the device with the other arrays: values that are not device arrays would be fixed when the training step is compiled.
+
# Let's train our network for a few epochs to see how well it performs on the training as well as test set:
-function custom_loss(model, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+function custom_loss(model, ps, st, (g, x, y, mask))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true), agg = l -> sum(l .* mask) / sum(mask))
st = Lux.trainmode(st)
- ŷ, st = model(g, x, ps, st)
- return logitcrossentropy(ŷ, y), st, 0
+ ŷ, st = model(g, x, ps, st)
+ return logitcrossentropy(ŷ, y), st, (;)
end
-function eval_loss_accuracy(model, ps, st, data_loader)
+function eval_loss_accuracy(model_compiled, ps, st, data_loader)
loss = 0.0
acc = 0.0
ntot = 0
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
for (g, y) in data_loader
- g = MLUtils.batch(g)
- n = length(y)
- ŷ, _ = model(g, g.ndata.x, ps, st)
- loss += logitcrossentropy(ŷ, y) * n
- acc += mean((ŷ .> 0) .== y) * n
+ g, y, _, n = pad_batch(g, y)
+ g = xdev(g)
+ ŷ, _ = model_compiled(g, g.x, ps, st)
+ ŷ, y = cdev(ŷ)[:, 1:n], y[:, 1:n]
+ loss += logitcrossentropy(ŷ, y) * n
+ acc += mean((ŷ .> 0) .== y) * n
ntot += n
end
return (loss = round(loss / ntot, digits = 4),
acc = round(acc * 100 / ntot, digits = 2))
end
-function train_model!(model, ps, st; epochs = 500, infotime = 100)
- train_state = Lux.Training.TrainState(model, ps, st, Adam(1e-2))
+function train_model(model, ps, st; epochs = 500, infotime = 100)
+ g, _, _, _ = pad_batch(first(train_loader)...)
+ g = xdev(g)
+ model_compiled = @compile model(g, g.x, ps, Lux.testmode(st))
- function report(epoch)
- train = eval_loss_accuracy(model, ps, st, train_loader)
+ function report(epoch, ps, st)
st = Lux.testmode(st)
- test = eval_loss_accuracy(model, ps, st, test_loader)
- st = Lux.trainmode(st)
+ train = eval_loss_accuracy(model_compiled, ps, st, train_loader)
+ test = eval_loss_accuracy(model_compiled, ps, st, test_loader)
@info (; epoch, train, test)
end
- report(0)
- for iter in 1:epochs
+
+ train_state = Training.TrainState(model, ps, st, Adam(1.0f-2))
+ report(0, ps, st)
+ for epoch in 1:epochs
for (g, y) in train_loader
- g = MLUtils.batch(g)
- _, loss, _, train_state = Lux.Training.single_train_step!(AutoZygote(), custom_loss,(g, g.ndata.x, y), train_state)
+ g, y, mask, _ = pad_batch(g, y)
+ g, y, mask = (g, y, mask) |> xdev
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss,
+ (g, g.x, y, mask), train_state)
end
- iter % infotime == 0 && report(iter)
+ epoch % infotime == 0 && report(epoch, train_state.parameters, train_state.states)
end
- return model, ps, st
+ return train_state.parameters, train_state.states
end
-model, ps, st = train_model!(model, ps, st);
+ps, st = train_model(model, ps, st);
-# As one can see, our model reaches around **74% test accuracy**.
+# As one can see, our model reaches around **70% test accuracy**.
# Reasons for the fluctuations in accuracy can be explained by the rather small dataset (only 38 test graphs), and usually disappear once one applies GNNs to larger datasets.
# ## (Optional) Exercise
diff --git a/GNNLux/docs/src_tutorials/node_classification.jl b/GNNLux/docs/src_tutorials/node_classification.jl
index 181a86601..02a4703da 100644
--- a/GNNLux/docs/src_tutorials/node_classification.jl
+++ b/GNNLux/docs/src_tutorials/node_classification.jl
@@ -12,11 +12,19 @@ using Lux, GNNLux
using MLDatasets
using Plots, TSne
using Random, Statistics
-using Zygote, Optimisers, OneHotArrays, ConcreteStructs
+using Reactant, Enzyme, Optimisers, OneHotArrays, ConcreteStructs
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
-rng = Random.seed!(17); # for reproducibility
+Random.seed!(17); # for reproducibility
+rng = Xoshiro(17); # for the model parameters and states
+
+# We compile the models with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate them with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+# Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+# We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+const xdev = reactant_device()
+const cdev = cpu_device();
# ## Visualize
# We want to visualize our results using t-distributed stochastic neighbor embedding (tsne) to project our output onto a 2D plane.
@@ -86,7 +94,7 @@ MLP = Chain(Dense(num_features => hidden_channels, relu),
Dropout(drop_rate),
Dense(hidden_channels => num_classes))
-ps, st = Lux.setup(rng, MLP);
+ps, st = Lux.setup(rng, MLP) |> xdev;
# ### Training a Multilayer Perceptron
@@ -96,37 +104,43 @@ ps, st = Lux.setup(rng, MLP);
# Let's train our simple MLP by following a similar procedure as described in [the first part of this tutorial](https://juliagraphs.org/GraphNeuralNetworks.jl/docs/GNNLux.jl/stable/tutorials/gnn_intro/).
# We again make use of the **cross entropy loss** and **Adam optimizer**.
# This time, we also define a **`accuracy` function** to evaluate how well our final model performs on the test node set (which labels have not been observed during training).
+# The training step is compiled by Reactant the first time `single_train_step!` is called, and the trained parameters and states are returned in the `train_state`.
+# The loss selects the training nodes by their integer indices `findall(train_mask)` rather than by the mask itself: Reactant currently gives wrong results when a boolean mask indexes the one-hot labels ([Reactant.jl#3087](https://github.com/EnzymeAD/Reactant.jl/issues/3087)).
+# For evaluation, we compile and run the forward pass in one go with `Reactant.@jit`, and compute the accuracy on the host.
-function loss(model, ps, st, x)
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
- ŷ, st = model(x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+function loss(model, ps, st, (x, y, idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
+ ŷ, st = model(x, ps, st)
+ return logitcrossentropy(ŷ[:, idx], y[:, idx]), st, (;)
end
-function train_model!(MLP, ps, st, x, epochs)
- train_state = Lux.Training.TrainState(MLP, ps, st, Adam(1e-3))
+function train_model(MLP, ps, st, data, epochs)
+ train_state = Training.TrainState(MLP, ps, st, Adam(1.0f-3))
for iter in 1:epochs
- _, loss_value, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss, x, train_state)
+ _, loss_value, _, train_state = Training.single_train_step!(AutoEnzyme(), loss, data, train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss_value)")
+ println("Epoch: $(iter) Loss: $(Float32(loss_value))")
end
end
+ return train_state.parameters, train_state.states
end
function accuracy(model, x, ps, st, y, mask)
st = Lux.testmode(st)
- ŷ, st = model(x, ps, st)
- mean(onecold(ŷ)[mask] .== onecold(y)[mask])
+ ŷ, _ = @jit model(x, ps, st)
+ mean(onecold(cdev(ŷ))[mask] .== onecold(y)[mask])
end
-train_model!(MLP, ps, st, x, 2000)
+x_dev, y_dev, train_idx = (x, y, findall(train_mask)) |> xdev
+
+ps, st = train_model(MLP, ps, st, (x_dev, y_dev, train_idx), 2000);
# After training the model, we can call the `accuracy` function to see how well our model performs on unseen labels.
# Here, we are interested in the accuracy of the model, *i.e.*, the ratio of correctly classified nodes:
-accuracy(MLP, x, ps, st, y, .!train_mask)
+accuracy(MLP, x_dev, ps, st, y, .!train_mask)
# As one can see, our MLP performs rather bad with only about ~50% test accuracy.
# But why does the MLP do not perform better?
@@ -184,9 +198,10 @@ end;
# Now let's visualize the node embeddings of our **untrained** GCN network.
gcn = GCN(num_features, num_classes, hidden_channels, drop_rate)
-ps, st = Lux.setup(rng, gcn)
-h_untrained, st = gcn(g, x, ps, st)
-h_untrained = h_untrained |> transpose
+ps, st = Lux.setup(rng, gcn) |> xdev
+g_dev = g |> xdev
+h_untrained, _ = @jit gcn(g_dev, x_dev, ps, Lux.testmode(st))
+h_untrained = h_untrained |> cdev |> transpose
visualize_tsne(h_untrained, g.ndata.targets)
@@ -195,38 +210,37 @@ visualize_tsne(h_untrained, g.ndata.targets)
-function loss(gcn, ps, st, tuple)
- g, x, y = tuple
- logitcrossentropy = CrossEntropyLoss(; logits=Val(true))
- ŷ, st = gcn(g, x, ps, st)
- return logitcrossentropy(ŷ[:, train_mask], y[:, train_mask]), (st), 0
+function loss(gcn, ps, st, (g, x, y, idx))
+ logitcrossentropy = CrossEntropyLoss(; logits = Val(true))
+ ŷ, st = gcn(g, x, ps, st)
+ return logitcrossentropy(ŷ[:, idx], y[:, idx]), st, (;)
end
-function train_model!(gcn, ps, st, g, x, y)
- train_state = Lux.Training.TrainState(gcn, ps, st, Adam(1e-2))
+function train_model(gcn, ps, st, data)
+ train_state = Training.TrainState(gcn, ps, st, Adam(1.0f-2))
for iter in 1:2000
- _, loss_value, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss,(g, x, y), train_state)
+ _, loss_value, _, train_state = Training.single_train_step!(AutoEnzyme(), loss, data, train_state)
if iter % 100 == 0
- println("Epoch: $(iter) Loss: $(loss_value)")
+ println("Epoch: $(iter) Loss: $(Float32(loss_value))")
end
end
- return gcn, ps, st
+ return train_state.parameters, train_state.states
end
-gcn, ps, st = train_model!(gcn, ps, st, g, x, y);
+ps, st = train_model(gcn, ps, st, (g_dev, x_dev, y_dev, train_idx));
# Now let's evaluate the loss of our trained GCN.
function accuracy(model, g, x, ps, st, y, mask)
st = Lux.testmode(st)
- ŷ, st = model(g, x, ps, st)
- mean(onecold(ŷ)[mask] .== onecold(y)[mask])
+ ŷ, _ = @jit model(g, x, ps, st)
+ mean(onecold(cdev(ŷ))[mask] .== onecold(y)[mask])
end
-train_accuracy = accuracy(gcn, g, g.ndata.features, ps, st, y, train_mask)
-test_accuracy = accuracy(gcn, g, g.ndata.features, ps, st, y, .!train_mask)
+train_accuracy = accuracy(gcn, g_dev, x_dev, ps, st, y, train_mask)
+test_accuracy = accuracy(gcn, g_dev, x_dev, ps, st, y, .!train_mask)
println("Train accuracy: $(train_accuracy)")
println("Test accuracy: $(test_accuracy)")
@@ -240,8 +254,8 @@ println("Test accuracy: $(test_accuracy)")
st = Lux.testmode(st) # inference mode
-out_trained, st = gcn(g, x, ps, st)
-out_trained = out_trained|> transpose
+out_trained, _ = @jit gcn(g_dev, x_dev, ps, st)
+out_trained = out_trained |> cdev |> transpose
visualize_tsne(out_trained, g.ndata.targets)
# ## (Optional) Exercises
diff --git a/GNNLux/docs/src_tutorials/traffic_prediction.jl b/GNNLux/docs/src_tutorials/traffic_prediction.jl
index 973455721..39fce102d 100644
--- a/GNNLux/docs/src_tutorials/traffic_prediction.jl
+++ b/GNNLux/docs/src_tutorials/traffic_prediction.jl
@@ -11,13 +11,20 @@
using Lux, GNNLux
using MLDatasets: METRLA
using Statistics, Plots, Random
-using Zygote, Optimisers
+using Reactant, Enzyme, Optimisers
const mae = MAELoss()
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true" # don't ask for dataset download confirmation
rng = Random.seed!(42); # for reproducibility
+# We compile the model with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) and differentiate it with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+# Reactant runs the compiled programs on a GPU if one is functional, and on the CPU otherwise.
+# We move data, graphs, parameters and states to the device with `xdev` and bring results back to the host with `cdev`:
+
+const xdev = reactant_device()
+const cdev = cpu_device();
+
# ## Dataset: METR-LA
# We use the `METR-LA` dataset from the paper [Diffusion Convolutional Recurrent Neural Network: Data-driven Traffic Forecasting](https://arxiv.org/pdf/1707.01926.pdf), which contains traffic data from loop detectors in the highway of Los Angeles County. The dataset contains traffic speed data from March 1, 2012 to June 30, 2012. The data is collected every 5 minutes, resulting in 12 observations per hour, from 207 sensors. Each sensor is a node in the graph, and the edge weights are the distances between the sensor locations.
@@ -64,10 +71,10 @@ plot_data(features[1:288],1) # Plot the speed of the first sensor for the first
# Now let's construct the static graph, the `train_loader` and `data_loader`.
-graph = GNNGraph(g.edge_index; edata = g.edge_data, g.num_nodes);
+graph = GNNGraph(g.edge_index; edata = g.edge_data, g.num_nodes) |> xdev;
-train_loader = zip(features[1:288], targets[1:288]); # train on 24 hours
-test_loader = zip(features[289:577], targets[289:577]); # test on next 24 hours
+train_loader = zip(features[1:288], targets[1:288]) |> collect |> xdev; # train on 24 hours
+test_loader = zip(features[289:577], targets[289:577]) |> collect |> xdev; # test on next 24 hours
# ## Model: T-GCN
@@ -78,10 +85,14 @@ model = GNNChain(TGCN(2 => 100; add_self_loops = false), Dense(100, 1))
# Let's look at the output of the model for the first batch of the training data.
# In Lux, the parameters and the state of the model are kept separate from the model itself, so we first initialize them.
-ps, st = Lux.setup(rng, model);
+ps, st = Lux.setup(rng, model) |> xdev;
+
+# All the inputs have the same shape, so we compile the forward pass once with `Reactant.@compile` and reuse it for all inputs:
+
+model_compiled = @compile model(graph, xdev(features[1]), ps, st)
-ŷ, _ = model(graph, features[1], ps, st);
-ŷ
+ŷ, _ = model_compiled(graph, xdev(features[1]), ps, st);
+cdev(ŷ)
# The output of the model is a tensor of size `(1, 3, 207)`, which corresponds to the dimension of the feature (in this case speed), the number of time steps, and the number of nodes in the graph, respectively. The model outputs the predicted traffic speed for each sensor at each time step.
@@ -90,6 +101,7 @@ ps, st = Lux.setup(rng, model);
# ## Training
# We train the model for 100 epochs, using the Adam optimizer with a learning rate of 0.001. We use the mean absolute error (MAE) as the loss function.
+# The first call to `single_train_step!` compiles the whole training step (forward pass, Enzyme gradient and Adam update), which is then reused for all the following steps.
function loss_function(model, ps, st, (graph, x, y))
ŷ, st = model(graph, x, ps, st)
@@ -98,16 +110,17 @@ end
function train(graph, train_loader, model, ps, st)
- train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001))
+ train_state = Training.TrainState(model, ps, st, Adam(0.001f0))
for epoch in 1:100
+ losses = Float32[]
for (x, y) in train_loader
- _, _, _, train_state = Lux.Training.single_train_step!(AutoZygote(), loss_function, (graph, x, y), train_state)
+ _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), loss_function, (graph, x, y), train_state)
+ push!(losses, loss)
end
if epoch % 10 == 0
- ps, st = train_state.parameters, train_state.states
- loss = mean([mae(model(graph, x, ps, st)[1], y) for (x, y) in train_loader])
+ loss = mean(losses)
@show epoch, loss
end
end
@@ -122,7 +135,7 @@ function plot_predicted_data(graph, features, targets, sensor)
ground_truth = []
for i in 1:3:length(features)
push!(ground_truth,targets[i][1,:,sensor])
- push!(prediction, model(graph, features[i], ps, st)[1][1,:,sensor])
+ push!(prediction, cdev(model_compiled(graph, xdev(features[i]), ps, st)[1])[1,:,sensor])
end
prediction = reduce(vcat,prediction)
ground_truth = reduce(vcat, ground_truth)
@@ -137,7 +150,7 @@ accuracy(ŷ, y) = 1 - Statistics.norm(y-ŷ)/Statistics.norm(y)
# Test accuracy:
-mean([accuracy(model(graph, x, ps, st)[1], y) for (x, y) in test_loader])
+mean([accuracy(cdev(model_compiled(graph, x, ps, st)[1]), cdev(y)) for (x, y) in test_loader])
# The accuracy is not very good but can be improved by training using more data. We used a small subset of the dataset for this tutorial because of the computational cost of training the model. From the plot of the predictions, we can see that the model is able to capture the general trend of the traffic speed, but it is not able to capture the peaks of the traffic.
diff --git a/GNNLux/test/Project.toml b/GNNLux/test/Project.toml
index 0cf4911ae..f457bc7fd 100644
--- a/GNNLux/test/Project.toml
+++ b/GNNLux/test/Project.toml
@@ -1,6 +1,7 @@
[deps]
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
GNNGraphs = "aed8fd31-079b-4b5a-b342-a13352159b8c"
GNNLux = "e8545f4d-a905-48ac-a8c4-ca114b98986d"
@@ -10,8 +11,10 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
diff --git a/GNNLux/test/reactant.jl b/GNNLux/test/reactant.jl
new file mode 100644
index 000000000..ab966888a
--- /dev/null
+++ b/GNNLux/test/reactant.jl
@@ -0,0 +1,98 @@
+@testitem "Reactant: forward pass" setup=[TestModuleLux] begin
+ using Reactant
+ Reactant.set_default_backend("cpu")
+ dev = reactant_device(; force = true)
+ rng = StableRNG(1234)
+
+ n, m = 10, 30
+ g = rand_graph(rng, n, m)
+ gw = rand_graph(rng, n, m, edge_weight = rand(rng, Float32, m ÷ 2))
+ x = randn(rng, Float32, 3, n)
+ ew = rand(rng, Float32, m)
+
+ # Each case exercises one of the `propagate` fast paths overridden by GNNlibReactantExt:
+ # `copy_xj`, `e_mul_xj` (explicit edge weights) and `w_mul_xj` (graph edge weights).
+ l = GCNConv(3 => 5, relu)
+ lw = GCNConv(3 => 5, relu, use_edge_weight = true)
+ ps, st = LuxCore.setup(rng, l)
+ for (layer, graph, args) in [(l, g, ()), (l, g, (ew,)), (lw, gw, ())]
+ y, _ = layer(graph, x, args..., ps, st)
+ # graph traced as a device input
+ yr, _ = @jit layer(dev(graph), dev(x), dev.(args)..., dev(ps), dev(st))
+ @test Array(yr) ≈ y
+ # graph kept on the host, i.e. embedded as a constant
+ yr, _ = @jit layer(graph, dev(x), dev.(args)..., dev(ps), dev(st))
+ @test Array(yr) ≈ y
+ end
+
+ # batched graph, with a traced graph indicator
+ gb = batch([rand_graph(rng, k, 2k, ndata = randn(rng, Float32, 3, k)) for k in 5:8])
+ model = GNNChain(GCNConv(3 => 5, relu), GlobalPool(mean), Dense(5 => 2))
+ ps, st = LuxCore.setup(rng, model)
+ y, _ = model(gb, gb.x, ps, st)
+ yr, _ = @jit model(dev(gb), dev(gb.x), dev(ps), dev(st))
+ @test Array(yr) ≈ y
+end
+
+@testitem "Reactant: gradients and training" setup=[TestModuleLux] begin
+ using Reactant, Enzyme, Optimisers, Zygote
+ Reactant.set_default_backend("cpu")
+ dev = reactant_device(; force = true)
+ rng = StableRNG(1234)
+
+ g = rand_graph(rng, 10, 30)
+ x = randn(rng, Float32, 3, 10)
+ model = GNNChain(GCNConv(3 => 5, relu), GCNConv(5 => 2))
+ ps, st = LuxCore.setup(rng, model)
+
+ loss(ps, g, x, st) = sum(abs2, first(model(g, x, ps, st)))
+ grad(ps, g, x, st) = Enzyme.gradient(Reverse, loss, ps, Const(g), Const(x), Const(st))[1]
+ gref = Zygote.gradient(loss, ps, g, x, st)[1]
+ gr = @jit grad(dev(ps), dev(g), dev(x), dev(st))
+ @test Array(gr.layers.layer_1.weight) ≈ gref.layers.layer_1.weight
+ @test Array(gr.layers.layer_2.bias) ≈ gref.layers.layer_2.bias
+
+ objective(model, ps, st, (g, x)) = (sum(abs2, first(model(g, x, ps, st))), st, (;))
+ function train(ts, data, nsteps)
+ losses = Float32[]
+ for _ in 1:nsteps
+ _, l, _, ts = Training.single_train_step!(AutoEnzyme(), objective, data, ts)
+ push!(losses, Float32(l))
+ end
+ return losses
+ end
+ ts = Training.TrainState(model, dev(ps), dev(st), Adam(1.0f-2))
+ losses = train(ts, dev((g, x)), 10)
+ @test losses[end] < losses[1]
+end
+
+@testitem "Reactant: padded mini-batches" setup=[TestModuleLux] begin
+ using Reactant
+ Reactant.set_default_backend("cpu")
+ dev = reactant_device(; force = true)
+ rng = StableRNG(1234)
+
+ model = GNNChain(GCNConv(3 => 5, relu), GlobalPool(mean), Dense(5 => 2))
+ ps, st = LuxCore.setup(rng, model)
+ # batches with different numbers of graphs, nodes and edges
+ batches = [batch([rand_graph(rng, k, 2k, ndata = randn(rng, Float32, 3, k)) for k in ks])
+ for ks in (5:8, 4:6, [9])]
+ pad(g) = pad_graph(g, 40, 80, 5)
+
+ # one compiled model gives the right result for every padded batch
+ gp = dev(pad(batches[1]))
+ compiled = @compile model(gp, gp.x, dev(ps), dev(st))
+ for g in batches
+ y, _ = model(g, g.x, ps, st)
+ gd = dev(pad(g))
+ yr, _ = compiled(gd, gd.x, dev(ps), dev(st))
+ @test Array(yr)[:, 1:g.num_graphs] ≈ y
+ end
+
+ # graphs of another size do not match the compiled shapes
+ for g in (batches[2], pad_graph(batches[1], 41, 80, 5), pad_graph(batches[1], 40, 82, 5),
+ pad_graph(batches[1], 40, 80, 6))
+ gd = dev(g)
+ @test_throws "RuntimeProgramInputMismatch" compiled(gd, gd.x, dev(ps), dev(st))
+ end
+end
diff --git a/GNNlib/Project.toml b/GNNlib/Project.toml
index a0dd0597e..d2a979fba 100644
--- a/GNNlib/Project.toml
+++ b/GNNlib/Project.toml
@@ -1,7 +1,7 @@
name = "GNNlib"
uuid = "a6a84749-d869-43f8-aacc-be26a1996e48"
authors = ["Carlo Lucibello and contributors"]
-version = "1.4.1"
+version = "1.5.0-DEV"
[workspace]
projects = ["test", "docs", "../GNNGraphs"]
@@ -21,24 +21,27 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
[extensions]
GNNlibAMDGPUExt = "AMDGPU"
GNNlibCUDAExt = "CUDA"
GNNlibMetalExt = "Metal"
GNNlibMooncakeExt = "Mooncake"
+GNNlibReactantExt = "Reactant"
[compat]
AMDGPU = "1, 2.7"
CUDA = "5, 6"
ChainRulesCore = "1.24"
DataStructures = "0.18, 0.19"
-GNNGraphs = "1.6"
+GNNGraphs = "1.7"
LinearAlgebra = "1"
Metal = "1.0"
MLUtils = "0.4"
Mooncake = "0.5"
NNlib = "0.9"
Random = "1"
+Reactant = "0.2"
Statistics = "1"
julia = "1.10"
diff --git a/GNNlib/ext/GNNlibReactantExt.jl b/GNNlib/ext/GNNlibReactantExt.jl
new file mode 100644
index 000000000..1e00f9733
--- /dev/null
+++ b/GNNlib/ext/GNNlibReactantExt.jl
@@ -0,0 +1,36 @@
+module GNNlibReactantExt
+
+using Reactant: AnyTracedRMatrix
+using GNNlib: GNNlib, propagate, copy_xj, e_mul_xj, w_mul_xj
+using GNNGraphs: GNNGraph, COO_T, get_edge_weight
+
+###### PROPAGATE SPECIALIZATIONS ####################
+
+# The generic fast paths multiply by a sparse adjacency matrix. Reactant cannot trace one whose
+# eltype is traced (scalar indexing) or that is built from traced edge indices (stack overflow).
+# Fall back to gather/scatter message passing, which lowers to XLA gather/scatter ops.
+
+## COPY_XJ
+
+function GNNlib.propagate(::typeof(copy_xj), g::GNNGraph{<:COO_T}, ::typeof(+),
+ xi, xj::AnyTracedRMatrix, e)
+ propagate((xi, xj, e) -> copy_xj(xi, xj, e), g, +, xi, xj, e)
+end
+
+## E_MUL_XJ
+
+function GNNlib.propagate(::typeof(e_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+),
+ xi, xj::AnyTracedRMatrix, e::AbstractVector)
+ propagate((xi, xj, e) -> e_mul_xj(xi, xj, e), g, +, xi, xj, e)
+end
+
+## W_MUL_XJ
+
+# The weights come from the graph here, so they have to be passed on explicitly:
+# `w_mul_xj` with `e = nothing` would return unweighted messages.
+function GNNlib.propagate(::typeof(w_mul_xj), g::GNNGraph{<:COO_T}, ::typeof(+),
+ xi, xj::AnyTracedRMatrix, e::Nothing)
+ propagate((xi, xj, e) -> w_mul_xj(xi, xj, e), g, +, xi, xj, get_edge_weight(g))
+end
+
+end #module
diff --git a/GNNlib/src/GNNlib.jl b/GNNlib/src/GNNlib.jl
index 8fe51c9a8..8bad759e0 100644
--- a/GNNlib/src/GNNlib.jl
+++ b/GNNlib/src/GNNlib.jl
@@ -10,7 +10,7 @@ using ChainRulesCore: @non_differentiable, ignore_derivatives
using GNNGraphs
using .GNNGraphs: COO_T, ADJMAT_T, SPARSE_T,
check_num_nodes, check_num_edges,
- _to_coo_graph,
+ _to_coo_graph, _scatter,
EType, NType # for heteroconvs
include("utils.jl")
diff --git a/GNNlib/src/utils.jl b/GNNlib/src/utils.jl
index 4967e151f..e2ec49211 100644
--- a/GNNlib/src/utils.jl
+++ b/GNNlib/src/utils.jl
@@ -12,7 +12,7 @@ See also: [`reduce_edges`](@ref).
function reduce_nodes(aggr, g::GNNGraph, x)
@assert size(x)[end] == g.num_nodes
indexes = graph_indicator(g)
- return NNlib.scatter(aggr, x, indexes)
+ return _scatter(aggr, x, indexes, g.num_graphs)
end
"""
@@ -35,7 +35,7 @@ given a heterogeneous graph `g`. The aggregation operator `aggr` can be `+`,
`mean`, `max`, or `min`.
"""
function reduce_nodes(aggr, node_type, g::GNNHeteroGraph, x)
- return NNlib.scatter(aggr, x[node_type], graph_indicator(g, node_type))
+ return _scatter(aggr, x[node_type], graph_indicator(g, node_type), g.num_graphs)
end
"""
@@ -49,7 +49,7 @@ function reduce_edges(aggr, g::GNNGraph, e)
@assert size(e)[end] == g.num_edges
s, t = edge_index(g)
indexes = graph_indicator(g)[s]
- return NNlib.scatter(aggr, e, indexes)
+ return _scatter(aggr, e, indexes, g.num_graphs)
end
"""
@@ -60,7 +60,7 @@ Graph-wise softmax of the node features `x`.
function softmax_nodes(g::GNNGraph, x)
@assert size(x)[end] == g.num_nodes
gi = graph_indicator(g)
- max_ = gather(scatter(max, x, gi), gi)
+ max_ = gather(_scatter(max, x, gi, g.num_graphs), gi)
num = exp.(x .- max_)
den = reduce_nodes(+, g, num)
den = gather(den, gi)
@@ -75,7 +75,7 @@ Graph-wise softmax of the edge features `e`.
function softmax_edges(g::GNNGraph, e)
@assert size(e)[end] == g.num_edges
gi = graph_indicator(g, edges = true)
- max_ = gather(scatter(max, e, gi), gi)
+ max_ = gather(_scatter(max, e, gi, g.num_graphs), gi)
num = exp.(e .- max_)
den = reduce_edges(+, g, num)
den = gather(den, gi)
@@ -104,9 +104,11 @@ function softmax_edge_neighbors(g::AbstractGNNGraph, e)
if isempty(t)
return zeros_like(e, eltype(e), (size(e)[1:end-1]..., 0))
end
- max_ = gather(scatter(max, e, t), t)
+ # An explicit destination size is needed when `t` is traced (e.g. by Reactant).
+ n = g isa GNNGraph ? g.num_nodes : maximum(t)
+ max_ = gather(_scatter(max, e, t, n), t)
num = exp.(e .- max_)
- den = gather(scatter(+, num, t), t)
+ den = gather(_scatter(+, num, t, n), t)
return num ./ den
end
diff --git a/GNNlib/test/utils.jl b/GNNlib/test/utils.jl
index 784329db9..5ebb20275 100644
--- a/GNNlib/test/utils.jl
+++ b/GNNlib/test/utils.jl
@@ -17,6 +17,14 @@
r2 = reduce_nodes(mean, graph_indicator(g), x)
@test r2 == r
+
+ # trailing graphs without nodes or edges still get an output column
+ gs = [rand_graph(4, 6, ndata = rand(Dx, 4), edata = rand(De, 6)),
+ GNNGraph(Int[], Int[]; num_nodes = 2, ndata = rand(Dx, 2), edata = rand(De, 0)),
+ GNNGraph(Int[], Int[]; num_nodes = 0, ndata = rand(Dx, 0), edata = rand(De, 0))]
+ gb = MLUtils.batch(gs)
+ @test size(reduce_nodes(+, gb, gb.ndata.x)) == (Dx, 3)
+ @test size(reduce_edges(+, gb, gb.edata.e)) == (De, 3)
end
@testset "reduce_nodes" begin
diff --git a/Project.toml b/Project.toml
index e9bb9caa6..87fb4a697 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,3 +1,3 @@
[workspace]
-projects = ["docs", "GNNGraphs", "GNNlib", "GraphNeuralNetworks", "GNNLux"]
+projects = ["GNNGraphs", "GNNlib", "GraphNeuralNetworks", "GNNLux"]