From 2b7e248ec7bad77cb1aaad2f2db4678c38af6fa5 Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Fri, 30 Oct 2026 22:53:48 +0100 Subject: [PATCH 1/5] Support compiling models with Reactant Fixes #710. - GNNGraphs accepts traced index arrays in a `GNNGraph`, and a new Reactant extension makes `degree` work inside a trace. - A new GNNlib Reactant extension makes `propagate` fall back to gather/scatter for the `copy_xj`, `e_mul_xj` and `w_mul_xj` fast paths, as on CUDA. - GNNlib's node and edge reductions and softmaxes pass explicit output sizes to their scatters, which Reactant needs for traced indices. This also fixes `reduce_nodes`/`reduce_edges` returning too few columns for batches that end in graphs with no nodes or edges. GNNLux gets tests for the forward pass under Reactant and for Enzyme gradients and training inside a compiled program. GNNlib goes to 1.5.0-DEV and requires GNNGraphs 1.6. The root `docs` project leaves the workspace: MultiDocumenter pins JSON 0.21, which holds Reactant at 0.2.217, a version that does not build on Julia 1.13. --- CHANGELOG.md | 14 ++++++ GNNGraphs/Project.toml | 5 +- GNNGraphs/ext/GNNGraphsReactantExt.jl | 19 ++++++++ GNNGraphs/src/abstracttypes.jl | 7 ++- GNNGraphs/src/gatherscatter.jl | 2 +- GNNLux/test/Project.toml | 3 ++ GNNLux/test/reactant.jl | 67 +++++++++++++++++++++++++++ GNNlib/Project.toml | 7 ++- GNNlib/ext/GNNlibReactantExt.jl | 36 ++++++++++++++ GNNlib/src/GNNlib.jl | 2 +- GNNlib/src/utils.jl | 16 ++++--- GNNlib/test/utils.jl | 8 ++++ Project.toml | 2 +- 13 files changed, 173 insertions(+), 15 deletions(-) create mode 100644 GNNGraphs/ext/GNNGraphsReactantExt.jl create mode 100644 GNNLux/test/reactant.jl create mode 100644 GNNlib/ext/GNNlibReactantExt.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index ffdbc7993..ff46222ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,19 @@ 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 ([#710]). + +## 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 ([#710]). + +**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 ([#710]). + ## GNNGraphs.jl 1.6.0 - 2026-11-04 **Added** @@ -244,6 +257,7 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#703]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/703 [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704 [#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707 +[#710]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/710 [#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 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/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/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/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..51cc30693 --- /dev/null +++ b/GNNLux/test/reactant.jl @@ -0,0 +1,67 @@ +@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 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"] From 31c713f1105ce22cf03f70126d4e82b0e4c82e5c Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Fri, 30 Oct 2026 22:53:48 +0100 Subject: [PATCH 2/5] Port the GNNLux tutorials to Reactant and Enzyme Train the tutorials with `Training.single_train_step!(AutoEnzyme(), ...)` on a `reactant_device()`, as Lux recommends, instead of Zygote, and regenerate their pages. --- GNNLux/docs/Project.toml | 7 +- .../tutorials/beginner_tutorials/gnn_intro.md | 108 +- .../graph_classification.md | 126 +- .../beginner_tutorials/node_classification.md | 5692 +---------------- .../traffic_prediction.md | 443 +- GNNLux/docs/src_tutorials/gnn_intro.jl | 69 +- .../src_tutorials/graph_classification.jl | 93 +- .../docs/src_tutorials/node_classification.jl | 84 +- .../docs/src_tutorials/traffic_prediction.jl | 39 +- 9 files changed, 639 insertions(+), 6022 deletions(-) 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).* ![cover](assets/cover_gnn_intro.png) @@ -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..10cf1deb9 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,80 @@ 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. -```` +### Compiling with fixed shapes -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. +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, so it cannot be fed to a training step compiled for another mini-batch. +The general solution is to pad every batch to the same size. Since MUTAG is small, here we simply batch all training graphs (and, separately, all test graphs) into one giant graph and perform full-batch training: + +````julia +g_train = MLUtils.batch(train_data[1]) |> xdev +y_train = train_data[2] |> xdev +g_test = MLUtils.batch(test_data[1]) |> xdev +y_test = test_data[2] |> xdev; +```` -Let's train our network for a few epochs to see how well it performs on the training as well as test set: +Let's train our network to see how well it performs on the training as well as test set. +For evaluation, we compile the forward pass in inference mode once for each of the two giant graphs: ````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)) + logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) st = Lux.trainmode(st) ŷ, st = model(g, x, ps, st) - return logitcrossentropy(ŷ, y), st, 0 + return logitcrossentropy(ŷ, y), st, (;) end -function eval_loss_accuracy(model, ps, st, data_loader) - loss = 0.0 - acc = 0.0 - ntot = 0 - 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 - ntot += n - end - return (loss = round(loss / ntot, digits = 4), - acc = round(acc * 100 / ntot, digits = 2)) +function eval_loss_accuracy(model_compiled, ps, st, g, y) + ŷ, _ = model_compiled(g, g.x, ps, st) + ŷ, y = cdev((ŷ, y)) + logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) + return (loss = round(logitcrossentropy(ŷ, y), digits = 4), + acc = round(mean((ŷ .> 0) .== y) * 100, 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) + st_test = Lux.testmode(st) + model_train = @compile model(g_train, g_train.x, ps, st_test) + model_test = @compile model(g_test, g_test.x, ps, st_test) - 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_train, ps, st, g_train, y_train) + test = eval_loss_accuracy(model_test, ps, st, g_test, y_test) @info (; epoch, train, test) end - report(0) - for iter 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) - end - - iter % infotime == 0 && report(iter) + + train_state = Training.TrainState(model, ps, st, Adam(1.0f-2)) + report(0, ps, st) + for epoch in 1:epochs + _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss, + (g_train, g_train.x, y_train), train_state) + + 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.6708f0, acc = 69.33), test = (loss = 0.6849f0, acc = 55.26)) +[ Info: (epoch = 100, train = (loss = 0.4564f0, acc = 77.33), test = (loss = 0.5876f0, acc = 73.68)) +[ Info: (epoch = 200, train = (loss = 0.4428f0, acc = 78.33), test = (loss = 0.5969f0, acc = 73.68)) +[ Info: (epoch = 300, train = (loss = 0.4002f0, acc = 85.33), test = (loss = 0.8212f0, acc = 71.05)) +[ Info: (epoch = 400, train = (loss = 0.3546f0, acc = 84.0), test = (loss = 0.7233f0, acc = 75.0)) +[ Info: (epoch = 500, train = (loss = 0.3374f0, acc = 85.33), test = (loss = 0.7399f0, acc = 76.32)) ```` -As one can see, our model reaches around **74% test accuracy**. +As one can see, our model reaches around **75% 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..c1c4808a9 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,65 @@ 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. -# Let's train our network for a few epochs to see how well it performs on the training as well as test set: +# ### Compiling with fixed shapes + +# 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, so it cannot be fed to a training step compiled for another mini-batch. +# The general solution is to pad every batch to the same size. Since MUTAG is small, here we simply batch all training graphs (and, separately, all test graphs) into one giant graph and perform full-batch training: + +g_train = MLUtils.batch(train_data[1]) |> xdev +y_train = train_data[2] |> xdev +g_test = MLUtils.batch(test_data[1]) |> xdev +y_test = test_data[2] |> xdev; -function custom_loss(model, ps, st, tuple) - g, x, y = tuple - logitcrossentropy = CrossEntropyLoss(; logits=Val(true)) +# Let's train our network to see how well it performs on the training as well as test set. +# For evaluation, we compile the forward pass in inference mode once for each of the two giant graphs: + +function custom_loss(model, ps, st, (g, x, y)) + logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) 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) - loss = 0.0 - acc = 0.0 - ntot = 0 - 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 - ntot += n - end - return (loss = round(loss / ntot, digits = 4), - acc = round(acc * 100 / ntot, digits = 2)) +function eval_loss_accuracy(model_compiled, ps, st, g, y) + ŷ, _ = model_compiled(g, g.x, ps, st) + ŷ, y = cdev((ŷ, y)) + logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) + return (loss = round(logitcrossentropy(ŷ, y), digits = 4), + acc = round(mean((ŷ .> 0) .== y) * 100, 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) + st_test = Lux.testmode(st) + model_train = @compile model(g_train, g_train.x, ps, st_test) + model_test = @compile model(g_test, g_test.x, ps, st_test) - 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_train, ps, st, g_train, y_train) + test = eval_loss_accuracy(model_test, ps, st, g_test, y_test) @info (; epoch, train, test) end - report(0) - for iter 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) - end - - iter % infotime == 0 && report(iter) + + train_state = Training.TrainState(model, ps, st, Adam(1.0f-2)) + report(0, ps, st) + for epoch in 1:epochs + _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss, + (g_train, g_train.x, y_train), train_state) + + 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 **75% 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. From 77399f710fb1b7fc2d3a0765e83521f16e476e8c Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Thu, 5 Nov 2026 12:32:35 +0100 Subject: [PATCH 3/5] Link the changelog entries to #730 --- CHANGELOG.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ff46222ca..5608f2f5a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,15 +13,15 @@ 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 ([#710]). +- 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]). ## 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 ([#710]). +- 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 ([#710]). +- `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 @@ -257,8 +257,8 @@ Lux implementations of the graph convolutional, pooling, and temporal layers [#703]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/703 [#704]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/704 [#707]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/pull/707 -[#710]: https://github.com/JuliaGraphs/GraphNeuralNetworks.jl/issues/710 [#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 From 64904d3926705e13d2fbacb4734d9c9fdf1a7edc Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Thu, 5 Nov 2026 21:23:24 +0100 Subject: [PATCH 4/5] Add pad_graph and restore mini-batches in the graph classification tutorial pad_graph pads a graph to fixed numbers of nodes, edges and graphs, so Reactant can compile a model or training step once and reuse it for every mini-batch. The graph classification tutorial goes back to master's mini-batch training, padding each batch and masking the padding graphs out of the loss. --- CHANGELOG.md | 1 + GNNGraphs/src/GNNGraphs.jl | 1 + GNNGraphs/src/transform.jl | 83 +++++++++++++++++ GNNGraphs/test/transform.jl | 55 ++++++++++++ .../graph_classification.md | 90 ++++++++++++------- .../src_tutorials/graph_classification.jl | 78 ++++++++++------ GNNLux/test/reactant.jl | 34 +++++++ 7 files changed, 284 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5608f2f5a..ceda678cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ Entries link to the pull request that introduced them. **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]). ## GNNlib.jl — Unreleased (towards 1.5.0) 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/transform.jl b/GNNGraphs/src/transform.jl index a89035de7..26e139794 100644 --- a/GNNGraphs/src/transform.jl +++ b/GNNGraphs/src/transform.jl @@ -564,6 +564,89 @@ function add_nodes(g::GNNGraph{<:COO_T}, n::Integer; ndata = (;)) ndata, g.edata, g.gdata) 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. + +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 + """ set_edge_weight(g::GNNGraph, w::AbstractVector) 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/src/tutorials/beginner_tutorials/graph_classification.md b/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md index 10cf1deb9..1701d3eff 100644 --- a/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md +++ b/GNNLux/docs/src/tutorials/beginner_tutorials/graph_classification.md @@ -219,55 +219,81 @@ 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. -### Compiling with fixed shapes +### 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, so it cannot be fed to a training step compiled for another mini-batch. -The general solution is to pad every batch to the same size. Since MUTAG is small, here we simply batch all training graphs (and, separately, all test graphs) into one giant graph and perform full-batch training: +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 -g_train = MLUtils.batch(train_data[1]) |> xdev -y_train = train_data[2] |> xdev -g_test = MLUtils.batch(test_data[1]) |> xdev -y_test = test_data[2] |> xdev; +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; ```` -Let's train our network to see how well it performs on the training as well as test set. -For evaluation, we compile the forward pass in inference mode once for each of the two giant graphs: +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, (g, x, y)) - 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, (;) + ŷ, st = model(g, x, ps, st) + return logitcrossentropy(ŷ, y), st, (;) end -function eval_loss_accuracy(model_compiled, ps, st, g, y) - ŷ, _ = model_compiled(g, g.x, ps, st) - ŷ, y = cdev((ŷ, y)) +function eval_loss_accuracy(model_compiled, ps, st, data_loader) + loss = 0.0 + acc = 0.0 + ntot = 0 logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) - return (loss = round(logitcrossentropy(ŷ, y), digits = 4), - acc = round(mean((ŷ .> 0) .== y) * 100, digits = 2)) + for (g, y) in data_loader + 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) - st_test = Lux.testmode(st) - model_train = @compile model(g_train, g_train.x, ps, st_test) - model_test = @compile model(g_test, g_test.x, ps, st_test) + g, _, _, _ = pad_batch(first(train_loader)...) + g = xdev(g) + model_compiled = @compile model(g, g.x, ps, Lux.testmode(st)) function report(epoch, ps, st) st = Lux.testmode(st) - train = eval_loss_accuracy(model_train, ps, st, g_train, y_train) - test = eval_loss_accuracy(model_test, ps, st, g_test, y_test) + 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 train_state = Training.TrainState(model, ps, st, Adam(1.0f-2)) report(0, ps, st) for epoch in 1:epochs - _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss, - (g_train, g_train.x, y_train), train_state) + for (g, y) in train_loader + 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 epoch % infotime == 0 && report(epoch, train_state.parameters, train_state.states) end @@ -278,16 +304,16 @@ ps, st = train_model(model, ps, st); ```` ```` -[ Info: (epoch = 0, train = (loss = 0.6708f0, acc = 69.33), test = (loss = 0.6849f0, acc = 55.26)) -[ Info: (epoch = 100, train = (loss = 0.4564f0, acc = 77.33), test = (loss = 0.5876f0, acc = 73.68)) -[ Info: (epoch = 200, train = (loss = 0.4428f0, acc = 78.33), test = (loss = 0.5969f0, acc = 73.68)) -[ Info: (epoch = 300, train = (loss = 0.4002f0, acc = 85.33), test = (loss = 0.8212f0, acc = 71.05)) -[ Info: (epoch = 400, train = (loss = 0.3546f0, acc = 84.0), test = (loss = 0.7233f0, acc = 75.0)) -[ Info: (epoch = 500, train = (loss = 0.3374f0, acc = 85.33), test = (loss = 0.7399f0, acc = 76.32)) +[ 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 **75% 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/graph_classification.jl b/GNNLux/docs/src_tutorials/graph_classification.jl index c1c4808a9..70785cadd 100644 --- a/GNNLux/docs/src_tutorials/graph_classification.jl +++ b/GNNLux/docs/src_tutorials/graph_classification.jl @@ -135,52 +135,78 @@ 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. -# ### Compiling with fixed shapes +# ### 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, so it cannot be fed to a training step compiled for another mini-batch. -# The general solution is to pad every batch to the same size. Since MUTAG is small, here we simply batch all training graphs (and, separately, all test graphs) into one giant graph and perform full-batch training: +# 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; -g_train = MLUtils.batch(train_data[1]) |> xdev -y_train = train_data[2] |> xdev -g_test = MLUtils.batch(test_data[1]) |> xdev -y_test = test_data[2] |> xdev; +# 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 to see how well it performs on the training as well as test set. -# For evaluation, we compile the forward pass in inference mode once for each of the two giant graphs: +# 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, (g, x, y)) - 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, (;) + ŷ, st = model(g, x, ps, st) + return logitcrossentropy(ŷ, y), st, (;) end -function eval_loss_accuracy(model_compiled, ps, st, g, y) - ŷ, _ = model_compiled(g, g.x, ps, st) - ŷ, y = cdev((ŷ, y)) +function eval_loss_accuracy(model_compiled, ps, st, data_loader) + loss = 0.0 + acc = 0.0 + ntot = 0 logitcrossentropy = CrossEntropyLoss(; logits = Val(true)) - return (loss = round(logitcrossentropy(ŷ, y), digits = 4), - acc = round(mean((ŷ .> 0) .== y) * 100, digits = 2)) + for (g, y) in data_loader + 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) - st_test = Lux.testmode(st) - model_train = @compile model(g_train, g_train.x, ps, st_test) - model_test = @compile model(g_test, g_test.x, ps, st_test) + g, _, _, _ = pad_batch(first(train_loader)...) + g = xdev(g) + model_compiled = @compile model(g, g.x, ps, Lux.testmode(st)) function report(epoch, ps, st) st = Lux.testmode(st) - train = eval_loss_accuracy(model_train, ps, st, g_train, y_train) - test = eval_loss_accuracy(model_test, ps, st, g_test, y_test) + 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 train_state = Training.TrainState(model, ps, st, Adam(1.0f-2)) report(0, ps, st) for epoch in 1:epochs - _, loss, _, train_state = Training.single_train_step!(AutoEnzyme(), custom_loss, - (g_train, g_train.x, y_train), train_state) + for (g, y) in train_loader + 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 epoch % infotime == 0 && report(epoch, train_state.parameters, train_state.states) end @@ -189,7 +215,7 @@ end ps, st = train_model(model, ps, st); -# As one can see, our model reaches around **75% 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/test/reactant.jl b/GNNLux/test/reactant.jl index 51cc30693..a2ada06de 100644 --- a/GNNLux/test/reactant.jl +++ b/GNNLux/test/reactant.jl @@ -65,3 +65,37 @@ end 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)) + gd = dev(g) + @test_throws "RuntimeProgramInputMismatch" compiled(gd, gd.x, dev(ps), dev(st)) + end + # the number of graphs is not a shape but a constant baked into the compiled model, + # so a different one goes unnoticed + gp6 = dev(pad_graph(batches[1], 40, 80, 6)) + @test_broken size(first(compiled(gp6, gp6.x, dev(ps), dev(st)))) == (2, 6) +end From 8bf7ffbd94a785295244bd4fddc46d256315ebaf Mon Sep 17 00:00:00 2001 From: Simone Carlo Surace Date: Thu, 5 Nov 2026 23:07:27 +0100 Subject: [PATCH 5/5] Add graph_sizes to GNNGraph graph_sizes holds the number of nodes in each graph of a batch. It duplicates graph_indicator, but its length is the number of graphs, which makes that number part of the array shapes a compiled Reactant function checks. A compiled model called on a batch with a different number of graphs now raises an error instead of silently returning wrong results. The positional constructor derives graph_sizes from graph_indicator, and transformations that keep the graph indicator pass graph_sizes through. --- CHANGELOG.md | 1 + GNNGraphs/src/convert.jl | 4 ++-- GNNGraphs/src/gnngraph.jl | 29 ++++++++++++++++++++--------- GNNGraphs/src/sampling.jl | 2 +- GNNGraphs/src/transform.jl | 30 ++++++++++++++++-------------- GNNGraphs/test/gnngraph.jl | 30 ++++++++++++++++++++++++++++++ GNNLux/test/reactant.jl | 7 ++----- 7 files changed, 72 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ceda678cc..1a57d24b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ Entries link to the pull request that introduced them. **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) 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/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 26e139794..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,7 @@ 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 """ @@ -576,7 +576,9 @@ 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. +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. @@ -660,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) @@ -1159,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/GNNLux/test/reactant.jl b/GNNLux/test/reactant.jl index a2ada06de..ab966888a 100644 --- a/GNNLux/test/reactant.jl +++ b/GNNLux/test/reactant.jl @@ -90,12 +90,9 @@ end 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)) + 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 - # the number of graphs is not a shape but a constant baked into the compiled model, - # so a different one goes unnoticed - gp6 = dev(pad_graph(batches[1], 40, 80, 6)) - @test_broken size(first(compiled(gp6, gp6.x, dev(ps), dev(st)))) == (2, 6) end