diff --git a/CHANGELOG.md b/CHANGELOG.md index ba2999f81..686eefa60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -99,6 +99,8 @@ breaking release). ### Added +- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `apply_parametric(model, x, μ, ps)` applies a `Chain` to one sample — a vector or a `(q, p)` pair of vectors — with one `μ`, or to a batch — a `2d × 1 × k` tensor or a `(q, p)` pair of `d × 1 × k` tensors — with a vector of `k` values of `μ`, sample by sample; the output has the shape of the input. The `Chain` receives the pair `(x, μ)`, and the internal `concatenate_array_with_parameters` puts the flattened `μ` below the state in the element type of the state, side by side for a matrix and a vector of `μ`, and repeated in every column for a matrix and one `μ`. `ParametricLoss()(model, ps, input, output, μ)` is the relative norm of the error of `apply_parametric`, and `ZygotePullback` takes the triple `(input, output, μ)` with one `μ` or a vector of them. `ParametricDataLoader(data, parameters)` takes a `2d × nt × m` array or a `(q, p)` pair of `d × nt × m` arrays of one array type with a vector of `m` named tuples, and `ParametricDataLoader(ensemble_solution)` takes an `EnsembleSolution` of `(q, p)` states with the parameters of its problem. A `Batch(b)` draws each pair of a start position and a trajectory once per epoch, as for `DataLoader`, and a minibatch is the state, the state one time step later and the `μ` of each sample. `train!` takes a `ParametricDataLoader` and trains on `ParametricLoss()` when no `loss` is given. A `Batch` with a sequence length or a prediction window other than one, a parameter vector whose length is not the number of trajectories, no trajectory, too few time steps for the batch, a `μ` whose element type is not the data's, and a `μ` whose keys, their order or the sizes of its leaves differ from the first `μ`'s raise an `ArgumentError`; a `(q, p)` pair of two array types is a `MethodError`, as for `DataLoader` (K20). The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there `concatenate_array_with_parameters` promoted a `Float32` state to `Float64`, rejected a matrix of more than one column with one `μ`, and the loader and the batch checked their counts with `@assert`. + - **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it. - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and, in the testset "_mul and tensor_tensor_mul on Metal match the host", `_mul` with a dense and a `SymmetricMatrix` weight from the left, a dense one from the right, and `tensor_tensor_mul` on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 and Julia 1.13.1, both testsets pass in a cold process: 4 of 4 assertions in the `PoissonTensor` testset and 8 of 8 in the "_mul and tensor_tensor_mul on Metal match the host" testset. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index f476dae34..8bf3625ab 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -129,6 +129,8 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. The same call with `q = A` gives a `DataLoader` with `input_dim == 4`. The docstring at `src/data_loader/data_loader.jl:205` states the one-type signature. Closing it means widening the three signatures to `NamedTupleOfArrays{T, (:q, :p)}`, and the docstring with them. + `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with + it. ## Upstream @@ -145,6 +147,47 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. The blocks are not test failures: both runs are green, and every testset of ea0c397 has the same count in both. +### K28 · AbstractNeuralNetworks' loss norm of a `(q, p)` named tuple promotes `Float32` to `Float64`. + +- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:73` +- **kind:** upstream +- **found:** 2026-10-25, by the critics of the parametric core +- **evidence:** + `_norm(dx::NamedTuple) = sum(map(_norm, values(dx))) / √length(dx)` divides by the square root of + an `Int`, which is a `Float64`. So `ParametricLoss` returns a `Float64` on `Float32` `(q, p)` + data, and `train!` on a `Float32` `(q, p)` `ParametricDataLoader` returns a `Vector{Float64}` + history. Array data keep `Float32`. + + ```julia + import AbstractNeuralNetworks + x = (q = rand(Float32, 2, 1, 3), p = rand(Float32, 2, 1, 3)) + typeof(AbstractNeuralNetworks._norm(x)) # Float64 + ``` + +### K30 · The gradient of AbstractNeuralNetworks' loss of a `(q, p)` named tuple infers as `Union{Nothing, NetworkParameters}`. + +- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:86` +- **kind:** upstream +- **found:** #357, review of 2026-10-26; the same inference shows on `origin/main` at `a664f5c9` +- **evidence:** + On Julia 1.13.2 and `a664f5c9`, `Base.return_types` of the pullback of + `ps -> _compute_loss((q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp))` gives + `Tuple{Union{Nothing, NetworkParameters}}` in `Float32` and `Float64`, and the same function + with `vcat` in place of the named tuples gives the concrete `NetworkParameters{T, (:L1,), …}`. So + `ZygotePullback(ParametricLoss())` on `(q, p)` data infers its gradient as that `Union`, for one + sample and for a batch, and a plain sum of squares over the `(q, p)` output of a batch infers + concretely. The run-time gradient is a `NetworkParameters`. + + ```julia + using GeometricMachineLearning, Zygote + import AbstractNeuralNetworks + ps = NetworkParameters((L1 = (W = rand(1, 1), b = rand(1)),)) + q, p, yq, yp = rand(1), rand(1), rand(1), rand(1) + v, pb = Zygote.pullback(ps -> AbstractNeuralNetworks._compute_loss( + (q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp)), ps) + Base.return_types(pb, (typeof(v),)) # Tuple{Union{Nothing, NetworkParameters}} + ``` + ### K4 · Metal.jl 1.11.1 needs a Metal device to precompile, so a sandboxed run with no Metal cache fails before the tests - **location:** `test/Project.toml:16` diff --git a/docs/src/data_loader/data_loader.md b/docs/src/data_loader/data_loader.md index 2cb8a7bd7..340f4fb6d 100644 --- a/docs/src/data_loader/data_loader.md +++ b/docs/src/data_loader/data_loader.md @@ -165,4 +165,5 @@ Batch(::Int, ::Int, ::Int) GeometricMachineLearning.number_of_batches GeometricMachineLearning.convert_input_and_batch_indices_to_array GeometricMachineLearning.accuracy +GeometricMachineLearning.ParametricDataLoader ``` diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index 256c4673c..10d0d15e1 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -46,6 +46,8 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math TransformerLoss AutoEncoderLoss ReducedLoss +GeometricMachineLearning.ParametricLoss +GeometricMachineLearning.apply_parametric projection_error reduction_error ``` diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 747610220..45532f68e 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -11,6 +11,12 @@ using AbstractNeuralNetworks # `NamedTuple` of layers is walked with `Base.map`, which needs nothing from anybody. import NeuralNetworkParameters: NetworkParameters using NeuralNetworkParameters: mapstorage +# `flatten` puts the parameters of a system into the input of a network that takes them +# (`src/parametric.jl`); `parameter_eltype` checks their element type against the data's, and +# `parameterlayout` their shape against each other's. `map_cotangent` and `storage_gradient` turn the +# cotangent of a parameter set into its gradient in the reverse rule of a batch. +using NeuralNetworkParameters: flatten, parameter_eltype, parameterlayout, map_cotangent, + storage_gradient using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries @@ -205,6 +211,7 @@ export solve!, encoder, decoder export iterate +include("parametric.jl") include("loss/losses.jl") export AbstractPullback, ZygotePullback @@ -214,8 +221,12 @@ export DataLoader export Batch, train! include("data_loader/tensor_assign.jl") include("data_loader/batch.jl") +include("data_loader/parametric_data_loader.jl") include("data_loader/train.jl") +# The training of a network that takes the parameters of the system beside its input. +public apply_parametric, ParametricLoss, ParametricDataLoader + include("reduced_system/reduced_system.jl") export HRedSys, reduction_error, projection_error, integrate_reduced_system, diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..8353d640e --- /dev/null +++ b/src/data_loader/parametric_data_loader.jl @@ -0,0 +1,110 @@ +@doc raw""" + ParametricDataLoader(data, parameters) + ParametricDataLoader(ensemble_solution) + +Make a loader of trajectories of a system together with the parameters ``\mu`` of the system that +each trajectory belongs to. + +`data` is a `2d × nt × m` array, or a `(q, p)` named tuple of two `d × nt × m` arrays of one +array type: `m` trajectories of `nt` time steps, with `q` above `p` where they are stacked. +`parameters` is a vector of `m` named tuples, the ``\mu`` of each trajectory. Every ``\mu`` has the +keys of the first, in the same order and at every level, and the same sizes of its leaves. The data +and every ``\mu`` have one element type. An `EnsembleSolution` of a problem with `(q, p)` states +gives its `q` and `p` and the parameters of its problem. + +# The batch + +A `Batch(batch_size)` draws, in one epoch, each pair of a start position and a trajectory once, in +minibatches of `batch_size`; the last minibatch is shorter. A drawn minibatch is the triple +`(input, output, μs)`: the state at the start position, the state one time step later, and the +``\mu`` of each sample. [`train!`](@ref) on a `ParametricDataLoader` trains on +[`ParametricLoss`](@ref)`()` when no `loss` is given, so the network is applied with +[`apply_parametric`](@ref), one time step per sample. A `Batch` with a sequence length or a +prediction window other than one is an `ArgumentError`. + +A problem without system parameters is trained with a [`DataLoader`](@ref). +""" +struct ParametricDataLoader{ + T, AT <: ArrayOrNamedTuple{T}, PT <: AbstractVector{<:NamedTuple}} + input::AT + parameters::PT + input_time_steps::Int + n_params::Int + + function ParametricDataLoader(input::AT, + parameters::PT) where {T, + AT <: Union{AbstractArray{T, 3}, + NamedTuple{(:q, :p), Tuple{AT3, AT3}} where {AT3 <: AbstractArray{T, 3}}}, + PT <: AbstractVector{<:NamedTuple}} + input_time_steps, n_params = data_size(input) + n_params ≥ 1 && length(parameters) == n_params || throw(ArgumentError( + "the data hold $(n_params) trajectories and there are $(length(parameters)) " * + "system parameters; a `ParametricDataLoader` takes one for each trajectory, " * + "and at least one trajectory.")) + layout = parameterlayout(first(parameters)) + for (j, μ) in enumerate(parameters) + S = parameter_eltype(μ) + S === T || S === Union{} || + throw(ArgumentError( + "the data have the element type $(T) and a system parameter has the " * + "element type $(S); a `ParametricDataLoader` takes one element type.")) + parameterlayout(μ) == layout || throw(ArgumentError( + "system parameter $(j) differs from system parameter 1 in its keys, their " * + "order or the sizes of its leaves; a `ParametricDataLoader` takes one shape.")) + end + new{T, AT, PT}(input, parameters, input_time_steps, n_params) + end +end + +# The number of time steps and the number of trajectories of the data. +data_size(input::AbstractArray{<:Any, 3}) = (size(input, 2), size(input, 3)) + +function data_size(input::NamedTuple{(:q, :p)}) + size(input.q) == size(input.p) || throw(ArgumentError( + "`q` has the size $(size(input.q)) and `p` has the size $(size(input.p)).")) + data_size(input.q) +end + +function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, + T1, + Vector{ST}}) where { + T, T1, TuT, TT <: TimeSeries{T1}, + ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}}} + ParametricDataLoader( + (q = data_tensor([s.q for s in ensemble_solution.s]), + p = data_tensor([s.p for s in ensemble_solution.s])), + ensemble_solution.problem.parameters) +end + +function networkbackend(dl::ParametricDataLoader{T, <:AbstractArray}) where {T} + networkbackend(dl.input) +end +function networkbackend(dl::ParametricDataLoader{T, <:NamedTuple}) where {T} + networkbackend(dl.input.q) +end + +# The start positions are those with a target `prediction_window` steps later; the network takes one +# time step per sample and returns one, so the sequence length and the prediction window are one. +function (batch::Batch)(dl::ParametricDataLoader) + batch.seq_length == 1 || throw(ArgumentError( + "a `ParametricDataLoader` draws one time step per sample, and the `Batch` asks for a " * + "sequence length of $(batch.seq_length).")) + batch.prediction_window == 1 || throw(ArgumentError( + "a `ParametricDataLoader` draws a target of one time step per sample, and the `Batch` " * + "asks for a prediction window of $(batch.prediction_window).")) + n = dl.input_time_steps - batch.prediction_window + n ≥ 1 || throw(ArgumentError( + "the `ParametricDataLoader` has $(dl.input_time_steps) time steps and the `Batch` asks " * + "for a prediction window of $(batch.prediction_window), which leaves no minibatch.")) + batch_over_two_axes(batch, n, dl.n_params) +end + +# The input at the start position, the target `seq_length` steps later, and the `μ` of the trajectory +# of each sample, which row 2 of the batch indices names. +function convert_input_and_batch_indices_to_array(dl::ParametricDataLoader, batch::Batch, + batch_indices_tuple::Vector{Tuple{Int, Int}}) + indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) + (gather(dl.input, indices, 0, batch.seq_length), + gather(dl.input, indices, batch.seq_length, batch.prediction_window), + [dl.parameters[trajectory] for (_, trajectory) in batch_indices_tuple]) +end diff --git a/src/data_loader/train.jl b/src/data_loader/train.jl index 2ced7f23a..b3583eef7 100644 --- a/src/data_loader/train.jl +++ b/src/data_loader/train.jl @@ -2,7 +2,8 @@ train!(nn, opt, dl, batch, n_epochs; loss, show_progress = true) Train the neural network `nn` with the `GeometricOptimizers.TrainingOptimizer` `opt` on the data of -the [`DataLoader`](@ref) `dl` for `n_epochs` epochs, and return the loss history. +the [`DataLoader`](@ref) or [`ParametricDataLoader`](@ref) `dl` for `n_epochs` epochs, and return +the loss history. Each epoch draws the minibatches of `dl` with the [`Batch`](@ref) `batch` and takes one step per minibatch: `loss` computes the loss on the minibatch and its gradient, and @@ -15,8 +16,9 @@ gradient, a zero leaf included. Entry `i` of the history is the mean of the mini trained through a [`ZygotePullback`](@ref). Without it, a [`TransformerIntegrator`](@ref) with a `Batch{:Transformer}` trains on [`TransformerLoss`](@ref)`(batch)`, a [`NeuralNetworkIntegrator`](@ref) with a `Batch{:FeedForward}` on `FeedForwardLoss()`, and an -autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`; any other combination -needs `loss`, and `train!` raises an `ArgumentError` without it. +autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`, and any network on a +`ParametricDataLoader` on [`ParametricLoss`](@ref)`()`; any other combination needs `loss`, and +`train!` raises an `ArgumentError` without it. A network of HamiltonianNeuralNetworks.jl trains on vector-field data with `loss = SymbolicPullback(arch)` of that package, and on consecutive states with a loss such as its @@ -29,9 +31,9 @@ with the same `opt` resumes the training where the first one stopped; the only r `train!` draws are those of the minibatch shuffle. `n_epochs < 1` is an `ArgumentError`, raised before any step. """ -function train!(nn::NeuralNetwork, opt::TrainingOptimizer, dl::DataLoader, batch::Batch, - n_epochs::Integer; - loss::Union{NetworkLoss, AbstractPullback} = NetworkLoss(nn.architecture, batch), +function train!(nn::NeuralNetwork, opt::TrainingOptimizer, + dl::Union{DataLoader, ParametricDataLoader}, batch::Batch, n_epochs::Integer; + loss::Union{NetworkLoss, AbstractPullback} = default_loss(nn, dl, batch), show_progress::Bool = true) n_epochs ≥ 1 || throw(ArgumentError("`train!` needs at least one epoch; `n_epochs` is $(n_epochs).")) @@ -66,6 +68,13 @@ end _as_pullback(l::NetworkLoss) = ZygotePullback(l) _as_pullback(p::AbstractPullback) = p +# A `DataLoader` takes the default loss of the architecture and the batch; a `ParametricDataLoader` +# draws the system parameters of each sample, which `ParametricLoss` hands to the network. +function default_loss(nn::NeuralNetwork, ::DataLoader, batch::Batch) + NetworkLoss(nn.architecture, batch) +end +default_loss(::NeuralNetwork, ::ParametricDataLoader, ::Batch) = ParametricLoss() + # The default loss of an architecture and a batch. The two-argument methods live here and not beside # `NetworkLoss(::AutoEncoder)` in `src/loss/losses.jl`, because `Batch` is included after `src/loss/`. function AbstractNeuralNetworks.NetworkLoss( diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 28ae38893..0820e5d28 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -226,3 +226,29 @@ function (loss::ReducedLoss)(model::Chain, ps::NetworkParameters, input::ArrayOrNamedTuple, output::ArrayOrNamedTuple) _compute_loss(loss.decoder(model(loss.encoder(input), ps)), output) end + +@doc raw""" + ParametricLoss() + +Make an instance of the loss of a network whose forward pass takes the parameters ``\mu`` of the +system beside its input. + +`ParametricLoss()(model, ps, input, output, μ)` applies the network with +[`apply_parametric`](@ref) and compares the result with `output` in the relative ``L_2`` norm: + +```math +\mathrm{loss}(\mathcal{NN}, \mathrm{input}, \mathrm{output}, \mu) = ||\mathcal{NN}(\mathrm{input}, \mu) - \mathrm{output}|| / ||\mathrm{output}||. +``` + +This is `FeedForwardLoss` with the system parameters passed to the network. `μ` is one +`OptionalParameters` for one sample, or a vector of them, one per sample of a batch, as a +[`ParametricDataLoader`](@ref) draws it. `train!` on a `ParametricDataLoader` trains on this loss +when no `loss` is given. +""" +struct ParametricLoss <: NetworkLoss end + +function (loss::ParametricLoss)(model::Chain, ps::NetworkParameters, + input::ArrayOrNamedTuple, output::ArrayOrNamedTuple, + μ::Union{OptionalParameters, AbstractVector}) + _compute_loss(apply_parametric(model, input, μ, ps), output) +end diff --git a/src/parametric.jl b/src/parametric.jl new file mode 100644 index 000000000..7767ce637 --- /dev/null +++ b/src/parametric.jl @@ -0,0 +1,139 @@ +# The forward pass of a network that takes the parameters of the *system* — not of the network — +# beside its input. The system parameters `μ` are an `OptionalParameters`: a `NamedTuple`, or +# `NullParameters()` for a system without any. + +# The flat vector of the system parameters, in the element type `T` where it is given, together with +# the `ParameterLayout` that puts such a vector back into the shape of `μ`. The layout is a value and +# not a closure, so a layer can store it in a field. `NullParameters` flattens to an empty vector. +_flatten_system_parameters(μ::NamedTuple) = flatten(μ) +_flatten_system_parameters(::NullParameters) = flatten(NamedTuple()) +_flatten_system_parameters(::Type{T}, μ::NamedTuple) where {T} = flatten(T, μ) +_flatten_system_parameters(::Type{T}, ::NullParameters) where {T} = flatten(T, NamedTuple()) + +# The input of a parametric layer: the state with the flattened system parameters below it, in the +# element type of the state. One `μ` goes below every column of a matrix; a vector of them goes one +# below each column, so the columns are concatenated side by side. Without system parameters the state +# is the input as it is. +concatenate_array_with_parameters(x::AbstractArray, ::NullParameters) = x + +function concatenate_array_with_parameters(x::AbstractVector, μ::NamedTuple) + vcat(x, first(_flatten_system_parameters(eltype(x), μ))) +end + +function concatenate_array_with_parameters(x::AbstractMatrix, μ::NamedTuple) + flat = first(_flatten_system_parameters(eltype(x), μ)) + vcat(x, repeat(flat, 1, size(x, 2))) +end + +function concatenate_array_with_parameters(x::AbstractMatrix, μs::AbstractVector) + size(x, 2) == length(μs) || throw(DimensionMismatch( + "the matrix has $(size(x, 2)) columns and there are $(length(μs)) system parameters")) + reduce(hcat, + [concatenate_array_with_parameters(x[:, i], μ) for (i, μ) in enumerate(μs)]) +end + +@doc raw""" + apply_parametric(model, x, μ, ps) + +Apply the `Chain` `model` with the network parameters `ps` to the input `x` and the parameters `μ` +of the system. + +`x` is one sample — a vector, or a `(q, p)` named tuple of vectors — with one `μ`, an +`OptionalParameters`. Or `x` is a batch of `k` samples — a `2d × 1 × k` tensor, or a `(q, p)` named +tuple of `d × 1 × k` tensors — with a vector of `k` values of `μ`, one per sample, as a +[`ParametricDataLoader`](@ref) draws them. A batch holds at least one sample. The output has the +shape of the input. + +The `Chain` receives the pair `(x, μ)`, so its first layer takes one sample and its `μ`, and a layer +that returns such a pair passes `μ` on to the next. A batch is applied sample by sample. + +# Examples + +```jldoctest +using GeometricMachineLearning +using GeometricMachineLearning: apply_parametric +import AbstractNeuralNetworks: AbstractExplicitLayer + +# a layer that adds the system parameter `a` to its input +struct AddParameter <: AbstractExplicitLayer{2, 2} end +(::AddParameter)((x, μ)::Tuple, ps) = x .+ μ.a + +model = Chain(AddParameter()) +ps = NetworkParameters((L1 = NamedTuple(),)) +x = reshape([1.0, 2.0, 3.0, 4.0], 2, 1, 2) +apply_parametric(model, x, [(a = 10.0,), (a = 20.0,)], ps) + +# output + +2×1×2 Array{Float64, 3}: +[:, :, 1] = + 11.0 + 12.0 + +[:, :, 2] = + 23.0 + 24.0 +``` +""" +function apply_parametric(model::Chain, + x::Union{ + AbstractVector, NamedTuple{(:q, :p), <:Tuple{AbstractVector, AbstractVector}}}, + μ::OptionalParameters, ps::NetworkParameters) + model((x, μ), ps) +end + +# A batch maps over the columns of its `d × k` reshape, not over indexed samples: under Zygote, the +# adjoint of each `getindex` allocates a cotangent of the whole batch, so indexed samples cost O(k²). +function apply_parametric(model::Chain, x::AbstractArray{<:Any, 3}, μs::AbstractVector, + ps::NetworkParameters) + check_parametric_batch(size(x), μs) + outputs = map_samples((xᵢ, μ, ps) -> apply_parametric(model, xᵢ, μ, ps), ps, + sample_columns(x), μs) + reshape(reduce(hcat, outputs), :, 1, length(μs)) +end + +# The output of a sample is stacked, `q` above `p`, so that its cotangent is an array, which the rule +# of `map_samples` hands to Zygote as it is. +function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, + μs::AbstractVector, ps::NetworkParameters) + check_parametric_batch(size(x.q), μs) + outputs = map_samples(ps, sample_columns(x.q), sample_columns(x.p), μs) do q, p, μ, ps + output = apply_parametric(model, (q = q, p = p), μ, ps) + vcat(output.q, output.p) + end + stacked = reduce(hcat, outputs) + d = size(x.q, 1) + (q = reshape(stacked[1:d, :], :, 1, length(μs)), + p = reshape(stacked[(d + 1):end, :], :, 1, length(μs))) +end + +# The samples of a `d × 1 × k` batch, as the `k` columns of its `d × k` reshape. +sample_columns(x::AbstractArray{<:Any, 3}) = eachcol(reshape(x, size(x, 1), size(x, 3))) + +# `f(sample..., ps)` for each sample, where sample `i` is element `i` of each of `samples`. +map_samples(f, ps, samples...) = map((sample...) -> f(sample..., ps), samples...) + +# Zygote's adjoint of `map` adds the gradients of the closure it maps, which holds `ps`, from `nothing`, +# so the gradient of `ps` infers as `Union{Nothing, NetworkParameters}`. This rule pulls each sample +# back on its own and adds the gradients of `ps`, as `NetworkParameters`, from the first sample on; a +# batch holds at least one sample. `f` holds the `Chain`, whose trainable parameters are in `ps`, so +# `f` has no gradient. +function ChainRulesCore.rrule(::typeof(map_samples), f, ps::NetworkParameters, samples...) + outputs_and_pullbacks = map((sample...) -> Zygote.pullback(f, sample..., ps), samples...) + function map_samples_pullback(Δoutputs) + Δs = map(((_, pullback), Δ) -> pullback(Δ), outputs_and_pullbacks, unthunk(Δoutputs)) + Δps = sum(Δ -> map_cotangent(storage_gradient, ps, last(Δ)), Δs) + Δsamples = ntuple(i -> map(Δ -> Δ[i], Δs), length(samples)) + (NoTangent(), NoTangent(), Δps, Δsamples...) + end + map(first, outputs_and_pullbacks), map_samples_pullback +end + +# A batch holds one time step per sample and one `μ` per sample. +function check_parametric_batch(input_size::NTuple{3, Int}, μs::AbstractVector) + input_size[2] == 1 && input_size[3] == length(μs) || throw(DimensionMismatch( + "a batch for `apply_parametric` has one time step and one system parameter per " * + "sample; the input has $(input_size[2]) time steps and $(input_size[3]) samples, " * + "and there are $(length(μs)) system parameters")) + nothing +end diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index 6848f7629..4377e0e5f 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -25,9 +25,10 @@ struct ZygotePullback{NNLT} <: AbstractPullback{NNLT} loss::NNLT end -# One method for a network that is trained on inputs only, and one for a network that is trained on -# an `(input, output)` pair: the two shapes a `NetworkLoss` takes. A method per shape and not a -# union over them, which says nothing about which one arrived. +# One method for a network that is trained on inputs only, one for a network that is trained on +# an `(input, output)` pair, and two for a network that also takes the parameters of the system, +# one `μ` or one per sample: the shapes a `NetworkLoss` takes. A method per shape and not a union +# over them, which says nothing about which one arrived. function (_pullback::ZygotePullback)(ps::NetworkParameters, model, input_nt::ArrayOrNamedTuple) value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps) value, Δ -> only(pb(Δ)) @@ -40,3 +41,18 @@ function (_pullback::ZygotePullback)( ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) value, Δ -> only(pb(Δ)) end + +function (_pullback::ZygotePullback)( + ps::NetworkParameters, model, + input_output_μ::Tuple{ + <:ArrayOrNamedTuple, <:ArrayOrNamedTuple, <:OptionalParameters}) + value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_output_μ...), ps) + value, Δ -> only(pb(Δ)) +end + +function (_pullback::ZygotePullback)( + ps::NetworkParameters, model, + input_output_μs::Tuple{<:ArrayOrNamedTuple, <:ArrayOrNamedTuple, <:AbstractVector}) + value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_output_μs...), ps) + value, Δ -> only(pb(Δ)) +end diff --git a/test/data_loader/parametric_data_loader.jl b/test/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..9557d0b30 --- /dev/null +++ b/test/data_loader/parametric_data_loader.jl @@ -0,0 +1,179 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ParametricDataLoader, ParametricLoss, params, + convert_input_and_batch_indices_to_array +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Test +import Random + +include("../helpers/eltypes.jl") +include("../helpers/parametric_affine.jl") + +# `m` trajectories of `nt` time steps of a state of dimension `2d`. +const d, nt, m = 1, 5, 3 + +# The data of trajectory `j` at time step `t`, in every entry, and its system parameters `μⱼ`. +constant_data(::Type{T}) where {T} = [T(j) for _ in 1:(2d), t in 1:nt, j in 1:m] +timed_data(::Type{T}) where {T} = [T(100j + t) for _ in 1:(2d), t in 1:nt, j in 1:m] +system_parameters(::Type{T}) where {T} = [(a = T(j),) for j in 1:m] + +# The loader of a `2d × nt × m` array, and of the `(q, p)` pair of its two halves. +function loaders(data, μs) + (ParametricDataLoader(data, μs), + ParametricDataLoader((q = data[1:d, :, :], p = data[(d + 1):end, :, :]), μs)) +end + +# Entry `k` of a drawn batch, as the vector of the state it holds. +sample(x::AbstractArray, k) = x[:, 1, k] +sample(x::NamedTuple, k) = vcat(x.q[:, 1, k], x.p[:, 1, k]) + +# Every sample of one epoch: its start position and trajectory, its input, its output and its `μ`. +function epoch(dl, batch) + samples = [] + for indices in batch(dl) + input, output, μs = convert_input_and_batch_indices_to_array(dl, batch, indices) + @test length(μs) == length(indices) + for (k, index) in enumerate(indices) + push!(samples, (index, sample(input, k), sample(output, k), μs[k])) + end + end + samples +end + +all_pairs = sort(vec([(t, j) for t in 1:(nt - 1), j in 1:m])) + +@testset "one epoch of Batch(b) on a ParametricDataLoader ($T)" for T in REAL_ELTYPES + Random.seed!(1) + for dl in loaders(constant_data(T), system_parameters(T)) + samples = epoch(dl, Batch(4)) + # each pair of a start position and a trajectory once + @test sort([index for (index, _, _, _) in samples]) == all_pairs + for (_, input, output, μ) in samples + @test μ.a == input[1] + @test all(==(input[1]), input) + end + # the last minibatch is shorter: 12 pairs in minibatches of 5 + @test length.(Batch(5)(dl)) == [5, 5, 2] + end + for dl in loaders(timed_data(T), system_parameters(T)) + samples = epoch(dl, Batch(4)) + @test sort([index for (index, _, _, _) in samples]) == all_pairs + for ((t, j), input, output, μ) in samples + @test input == fill(T(100j + t), 2d) + # the output is the input one time step later + @test output == input .+ 1 + @test μ.a == T(j) + @test eltype(input) === eltype(output) === T + end + end +end + +@testset "a ParametricDataLoader of an EnsembleSolution" begin + parameters = [default_parameters(), map(v -> v .+ 1, default_parameters())] + solution = integrate(hodeensemble(; parameters = parameters), ImplicitMidpoint()) + dl = ParametricDataLoader(solution) + @test dl.parameters == parameters + @test dl.n_params == 2 + n_steps = length(solution.t) + @test dl.input_time_steps == n_steps + for j in 1:2, t in 1:n_steps + + @test dl.input.q[:, t, j] == solution.s[j].q[t - 1] + @test dl.input.p[:, t, j] == solution.s[j].p[t - 1] + end + Random.seed!(2) + samples = epoch(dl, Batch(7)) + @test sort([index for (index, _, _, _) in samples]) == + sort(vec([(t, j) for t in 1:(n_steps - 1), j in 1:2])) + for ((t, j), input, output, μ) in samples + @test μ == parameters[j] + @test input == vcat(dl.input.q[:, t, j], dl.input.p[:, t, j]) + @test output == vcat(dl.input.q[:, t + 1, j], dl.input.p[:, t + 1, j]) + end +end + +# The message of the `ArgumentError` that `f` throws. +function argument_error_message(f) + try + f() + catch error + error isa ArgumentError && return error.msg + rethrow() + end + "no error" +end + +@testset "the edges of a ParametricDataLoader" begin + data = timed_data(Float64) + μs = system_parameters(Float64) + dl = ParametricDataLoader(data, μs) + + # one time step per sample: a sequence length other than one is an error before any step + message = argument_error_message(() -> Batch(2, 3, 1)(dl)) + @test occursin("sequence length of 3", message) + # a prediction window other than one is an error before any step, with the time steps to spare + message = argument_error_message(() -> Batch{:FeedForward}(2, 1, 2)(dl)) + @test occursin("prediction window of 2", message) + + # a parameter vector of the wrong length, and no trajectory + message = argument_error_message(() -> ParametricDataLoader(data, μs[1:2])) + @test occursin("hold 3 trajectories", message) && + occursin("there are 2 system parameters", message) + message = argument_error_message( + () -> ParametricDataLoader(zeros(2d, nt, 0), NamedTuple{(:a,), Tuple{Float64}}[])) + @test occursin("hold 0 trajectories", message) + + # too few time steps for the batch + message = argument_error_message(() -> Batch(2)(ParametricDataLoader(data[:, 1:1, :], μs))) + @test occursin("1 time step", message) + + # one element type for the data and every `μ` + message = argument_error_message(() -> ParametricDataLoader(Float32.(data), μs)) + @test occursin("Float32", message) && occursin("Float64", message) + @test ParametricDataLoader(data, [(a = 1.0, b = [2.0, 3.0]) for _ in 1:m]).n_params == m + # an empty `μ` has no element type, and concatenates nothing + @test ParametricDataLoader(data, [(;) for _ in 1:m]).n_params == m + + # `q` and `p` of one size + message = argument_error_message( + () -> ParametricDataLoader((q = data[1:1, :, :], p = data[:, :, :]), μs)) + @test occursin("(1, $(nt), $(m))", message) && occursin("(2, $(nt), $(m))", message) + + # `q` and `p` of one array type, as `gather` batches them + @test_throws MethodError ParametricDataLoader( + (q = data[1:d, :, :], p = view(data, (d + 1):(2d), :, :)), μs) + views = ParametricDataLoader( + (q = view(data, 1:d, :, :), p = view(data, (d + 1):(2d), :, :)), μs) + input, output, _ = convert_input_and_batch_indices_to_array(views, Batch(2), [(1, 2)]) + @test sample(input, 1) == data[:, 1, 2] && sample(output, 1) == data[:, 2, 2] + + # every `μ` with the keys of the first, in its order and at every level, and its leaf sizes + for parameters in ([(a = 1.0,), (b = 2.0,), (a = 3.0,)], + [(a = 1.0, b = 2.0), (b = 3.0, a = 4.0), (a = 5.0, b = 6.0)], + [(a = 1.0,), (a = 2.0, b = 3.0), (a = 3.0,)], + [(a = (c = 1.0,),), (a = (e = 2.0,),), (a = (c = 3.0,),)], + [(a = [1.0, 2.0],), (a = [1.0, 2.0, 3.0],), (a = [1.0, 2.0],)]) + message = argument_error_message(() -> ParametricDataLoader(data, parameters)) + @test occursin("system parameter 2", message) + end +end + +# The network is `(x, μ) ↦ W [x; flatten(μ)] + b`. `train!` without `loss` and with +# `loss = ParametricLoss()` draw the same network and the same minibatches under one seed. +@testset "train! on a ParametricDataLoader trains on ParametricLoss ($T)" for T in REAL_ELTYPES + dl = ParametricDataLoader(timed_data(T) ./ 100, system_parameters(T)) + function trained(; kwargs...) + Random.seed!(3) + nn = NeuralNetwork(Chain(ParametricAffine{2d}(1)), T) + initial = deepcopy(params(nn)) + opt = TrainingOptimizer(nn.params; algorithm = Adam()) + history = train!(nn, opt, dl, Batch(4), 3; show_progress = false, kwargs...) + history, initial, params(nn) + end + history, initial, final = trained() + @test history == first(trained(; loss = ParametricLoss())) + @test length(history) == 3 + @test eltype(history) === T + @test final.L1.W != initial.L1.W + @test final.L1.b != initial.L1.b +end diff --git a/test/helpers/parametric_affine.jl b/test/helpers/parametric_affine.jl new file mode 100644 index 000000000..2fac1b5fb --- /dev/null +++ b/test/helpers/parametric_affine.jl @@ -0,0 +1,35 @@ +# A layer whose forward pass takes the parameters `μ` of the system beside the state: +# `(x, μ) ↦ W [x; flatten(μ)] + b`. It acts on one sample, a vector `x` or a `(q, p)` pair of vectors +# that stacks `q` above `p`, which is what `apply_parametric` hands the `Chain`. The tests build the +# expected value of this map by hand, with the flattened `μ` written out. +using GeometricMachineLearning: concatenate_array_with_parameters +using AbstractNeuralNetworks: AbstractExplicitLayer, Initializer, NamedTupleOfArrays +import AbstractNeuralNetworks: initialparameters, parameterlength +import KernelAbstractions +import Random + +struct ParametricAffine{M} <: AbstractExplicitLayer{M, M} + parameter_length::Int +end + +function initialparameters( + rng::Random.AbstractRNG, init::Initializer, l::ParametricAffine{M}, + backend::KernelAbstractions.Backend, ::Type{T}) where {M, T} + W = KernelAbstractions.zeros(backend, T, M, M + l.parameter_length) + b = KernelAbstractions.zeros(backend, T, M) + init(rng, W) + init(rng, b) + (W = W, b = b) +end + +parameterlength(l::ParametricAffine{M}) where {M} = M * (M + l.parameter_length) + M + +(l::ParametricAffine)(input::Tuple, ps::NamedTuple) = affine(input..., ps) + +affine(x::AbstractVector, μ, ps) = ps.W * concatenate_array_with_parameters(x, μ) + ps.b + +function affine(x::NamedTupleOfArrays{<:Any, (:q, :p)}, μ, ps) + y = affine(vcat(x.q, x.p), μ, ps) + n = length(x.q) + (q = y[1:n], p = y[(n + 1):end]) +end diff --git a/test/loss/parametric_loss.jl b/test/loss/parametric_loss.jl new file mode 100644 index 000000000..07eee3f16 --- /dev/null +++ b/test/loss/parametric_loss.jl @@ -0,0 +1,82 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ParametricLoss, ZygotePullback +using FiniteDifferences: central_fdm, grad +using LinearAlgebra: norm +using Test +import Random + +include("../helpers/eltypes.jl") +include("../helpers/parametric_affine.jl") + +# System parameters in `Float64`, with a vector-valued leaf, and their flattened form written out by +# hand in the element type `T`. +μ(j) = (a = Float64(j), b = [0.5j, -1.5j]) +flat(::Type{T}, μ) where {T} = T[μ.a; μ.b[1]; μ.b[2]] + +const K = 4 +const μS = [μ(j) for j in 1:K] + +# The network `(x, μ) ↦ W [x; flatten(μ)] + b` on a batch of `K` samples, and the relative norm of its +# error, both written out with the flattened `μ` built by hand. +function closed_form(W, b, X, μs) + T = eltype(X) + reshape(reduce(hcat, [W * vcat(X[:, 1, i], flat(T, μs[i])) + b for i in axes(X, 3)]), + size(X)) +end +closed_form_loss(W, b, X, Y, μs) = norm(closed_form(W, b, X, μs) - Y) / norm(Y) + +function data(::Type{T}) where {T} + rng = Random.Xoshiro(7) + W = randn(rng, T, 2, 5) + b = randn(rng, T, 2) + X = randn(rng, T, 2, 1, K) + Y = randn(rng, T, 2, 1, K) + W, b, X, Y +end + +network_parameters(W, b) = NetworkParameters((L1 = (W = W, b = b),)) +const MODEL = Chain(ParametricAffine{2}(3)) + +@testset "ParametricLoss is the relative norm of the error ($T)" for T in REAL_ELTYPES + W, b, X, Y = data(T) + value = ParametricLoss()(MODEL, network_parameters(W, b), X, Y, μS) + @test value isa T + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(T)) +end + +@testset "ZygotePullback of ParametricLoss against finite differences" begin + W, b, X, Y = data(Float64) + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (X, Y, μS)) + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(Float64)) + + gradient = back(one(value)) + # the finite-difference gradient of the closed form over the entries of `W`, then of `b` + loss_of_entries(v) = closed_form_loss(reshape(v[1:10], 2, 5), v[11:12], X, Y, μS) + reference = only(grad(central_fdm(5, 1), loss_of_entries, vcat(vec(W), b))) + computed = vcat(vec(gradient.L1.W), gradient.L1.b) + @test norm(computed - reference) ≤ 1e-6 * norm(reference) +end + +# One sample with one `μ` is the other shape the pullback takes. +@testset "ZygotePullback of ParametricLoss on one sample ($T)" for T in REAL_ELTYPES + W, b, X, Y = data(T) + x, y = X[:, 1, 1], Y[:, 1, 1] + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (x, y, μ(2))) + expected = norm(W * vcat(x, flat(T, μ(2))) + b - y) / norm(y) + @test value isa T + @test isapprox(value, expected; rtol = 10 * eps(T)) + @test eltype(back(one(value)).L1.W) === T +end + +@testset "ZygotePullback of ParametricLoss keeps Float32" begin + W, b, X, Y = data(Float32) + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (X, Y, μS)) + @test value isa Float32 + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(Float32)) + gradient = back(one(value)) + @test eltype(gradient.L1.W) === Float32 + @test eltype(gradient.L1.b) === Float32 +end diff --git a/test/parametric.jl b/test/parametric.jl new file mode 100644 index 000000000..77d5ce3fa --- /dev/null +++ b/test/parametric.jl @@ -0,0 +1,176 @@ +using GeometricMachineLearning +using GeometricMachineLearning: concatenate_array_with_parameters, apply_parametric, + NullParameters, ParametricLoss, ZygotePullback +using FiniteDifferences: central_fdm, grad +using LinearAlgebra: norm +using Test +import Random +import Zygote + +include("helpers/eltypes.jl") +include("helpers/parametric_affine.jl") + +# System parameters in `Float64`, with a vector-valued leaf; `flat` writes out by hand what +# flattening one of them gives in the element type `T`: the leaves in the order of the keys. +μ(j) = (a = Float64(j), b = [0.5j, -1.5j]) +flat(::Type{T}, μ) where {T} = T[μ.a; μ.b[1]; μ.b[2]] + +@testset "concatenate_array_with_parameters ($T)" for T in REAL_ELTYPES + x = T[1 2 3; 4 5 6] + + # one `μ` per column: the columns side by side + concatenated = concatenate_array_with_parameters(x, [μ(1), μ(2), μ(3)]) + @test concatenated == T[1 2 3; 4 5 6; 1 2 3; 0.5 1 1.5; -1.5 -3 -4.5] + @test eltype(concatenated) === T + + # one `μ` below a vector + concatenated = concatenate_array_with_parameters(x[:, 2], μ(2)) + @test concatenated == T[2, 5, 2, 1, -3] + @test eltype(concatenated) === T + + # one `μ` below every column of a matrix + concatenated = concatenate_array_with_parameters(x, μ(2)) + @test concatenated == T[1 2 3; 4 5 6; 2 2 2; 1 1 1; -3 -3 -3] + @test eltype(concatenated) === T + + # no system parameters: the state itself + @test concatenate_array_with_parameters(x, NullParameters()) === x + @test concatenate_array_with_parameters(x[:, 1], NullParameters()) == x[:, 1] + # an empty named tuple concatenates nothing + @test concatenate_array_with_parameters(x[:, 1], (;)) == x[:, 1] + @test eltype(concatenate_array_with_parameters(x[:, 1], (;))) === T + + @test_throws DimensionMismatch concatenate_array_with_parameters(x, [μ(1), μ(2)]) +end + +# `ParametricAffine` maps `(x, μ)` to `W [x; flatten(μ)] + b`; `closed_form` writes that out with the +# flattened `μ` built by hand. +@testset "apply_parametric against the closed form ($T)" for T in REAL_ELTYPES + W = T[0.5 -1 0.25 2 -0.5; 1.5 0.75 -2 0.125 1] + b = T[0.1, -0.3] + ps = NetworkParameters((L1 = (W = W, b = b),)) + model = Chain(ParametricAffine{2}(3)) + closed_form(x, μ) = W * vcat(x, flat(T, μ)) + b + tolerance = 10 * eps(T) + + # one sample: a vector, and a `(q, p)` pair of vectors + x = T[0.3, -0.7] + output = apply_parametric(model, x, μ(2), ps) + @test size(output) == size(x) + @test eltype(output) === T + @test isapprox(output, closed_form(x, μ(2)); rtol = tolerance) + + output = apply_parametric(model, (q = x[1:1], p = x[2:2]), μ(2), ps) + @test size(output.q) == size(output.p) == (1,) + @test eltype(output.q) === eltype(output.p) === T + @test isapprox(vcat(output.q, output.p), closed_form(x, μ(2)); rtol = tolerance) + + # a batch of `k` samples with `k` values of `μ`: a `2d × 1 × k` tensor, and a `(q, p)` pair of + # `d × 1 × k` tensors + k = 4 + μs = [μ(j) for j in 1:k] + X = reshape(T[0.3, -0.7, 1.1, 0.2, -0.4, 0.9, 0.0, 2.5], 2, 1, k) + expected = reshape(reduce(hcat, [closed_form(X[:, 1, i], μs[i]) for i in 1:k]), 2, 1, k) + + output = apply_parametric(model, X, μs, ps) + @test size(output) == size(X) + @test eltype(output) === T + @test isapprox(output, expected; rtol = tolerance) + + output = apply_parametric(model, (q = X[1:1, :, :], p = X[2:2, :, :]), μs, ps) + @test size(output.q) == size(output.p) == (1, 1, k) + @test eltype(output.q) === eltype(output.p) === T + @test isapprox(output.q, expected[1:1, :, :]; rtol = tolerance) + @test isapprox(output.p, expected[2:2, :, :]; rtol = tolerance) + + # one `μ` per sample and one time step per sample + @test_throws DimensionMismatch apply_parametric(model, X, μs[1:3], ps) + @test_throws DimensionMismatch apply_parametric( + model, reshape(X, 2, 2, 2), μs[1:2], ps) +end + +# The one-sample method takes a vector or a `(q, p)` pair of vectors; a pair of tensors with one `μ` +# matches no method. +@testset "apply_parametric takes one sample as vectors" begin + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + X = reshape([0.3, -0.7, 1.1, 0.2], 2, 1, 2) + @test applicable(apply_parametric, model, (q = [0.3], p = [-0.7]), (a = 2.0,), ps) + @test !applicable( + apply_parametric, model, (q = X[1:1, :, :], p = X[2:2, :, :]), (a = 2.0,), ps) +end + +# The bytes that the pullback of `ParametricLoss` and the call of its closure allocate on `data`: the +# least of three readings after two calls that compile. +function pullback_bytes(model, ps, data) + function run() + value, back = ZygotePullback(ParametricLoss())(ps, model, data) + back(one(value)) + end + run() + run() + minimum(@allocated(run()) for _ in 1:3) +end + +# A batch of `k` samples costs O(k): the bytes of one more sample, between 64 and 1024 samples, are +# at most twice the bytes per sample at 64. A sample indexed one at a time in the differentiated path +# costs a cotangent of the whole batch, O(k²) in all. +@testset "the pullback of a batch allocates linearly in its size ($form)" for form in (:array, :qp) + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + function batch(k) + rng = Random.Xoshiro(k) + X, Y = randn(rng, 2, 1, k), randn(rng, 2, 1, k) + μs = [(a = randn(rng),) for _ in 1:k] + form === :array ? (X, Y, μs) : + ((q = X[1:1, :, :], p = X[2:2, :, :]), (q = Y[1:1, :, :], p = Y[2:2, :, :]), μs) + end + b64 = pullback_bytes(model, ps, batch(64)) + b1024 = pullback_bytes(model, ps, batch(1024)) + @test (b1024 - b64) / 960 ≤ 2 * b64 / 64 +end + +# The gradient of the loss with respect to the input, against central finite differences of the same +# loss over the entries of the input. +@testset "the input gradient of ParametricLoss ($form)" for form in (:array, :qp) + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + rng = Random.Xoshiro(3) + k = 4 + X, Y = randn(rng, 2, 1, k), randn(rng, 2, 1, k) + μs = [(a = randn(rng),) for _ in 1:k] + as_input(v) = form === :array ? reshape(v, 2, 1, k) : + (q = reshape(v[1:k], 1, 1, k), p = reshape(v[(k + 1):end], 1, 1, k)) + as_vector(x) = form === :array ? vec(x) : vcat(vec(x.q), vec(x.p)) + y = form === :array ? Y : (q = Y[1:1, :, :], p = Y[2:2, :, :]) + loss(x) = ParametricLoss()(model, ps, x, y, μs) + v = form === :array ? vec(X) : vcat(vec(X[1, 1, :]), vec(X[2, 1, :])) + + computed = only(Zygote.gradient(loss, as_input(v))) + @test computed !== nothing + reference = only(grad(central_fdm(5, 1), loss ∘ as_input, v)) + @test norm(as_vector(computed) - reference) ≤ 1e-6 * norm(reference) +end + +# The gradient that the pullback of `ParametricLoss` returns infers as its run-time type, for one sample +# and for a batch. The `(q, p)` form is left out: AbstractNeuralNetworks' loss of a `(q, p)` named tuple +# infers its gradient as a `Union` on its own (KNOWN_ISSUES K30). +@testset "the gradient of ParametricLoss infers concretely ($T)" for T in REAL_ELTYPES + ps = NetworkParameters((L1 = (W = T[0.5 -1 0.25; 1.5 0.75 -2], b = T[0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + rng = Random.Xoshiro(5) + k = 4 + X, Y = randn(rng, T, 2, 1, k), randn(rng, T, 2, 1, k) + μs = [(a = randn(rng, T),) for _ in 1:k] + for data in ((X[:, 1, 1], Y[:, 1, 1], μs[1]), (X, Y, μs)) + value, back = ZygotePullback(ParametricLoss())(ps, model, data) + @test (@inferred back(one(value))) isa NetworkParameters{T} + end +end + +@testset "the parametric names are public and not exported" begin + for name in (:apply_parametric, :ParametricLoss, :ParametricDataLoader) + @test Base.ispublic(GeometricMachineLearning, name) + @test !Base.isexported(GeometricMachineLearning, name) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 8a25eef3d..a7e0b11ac 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,6 +70,8 @@ if "core" in GROUPS @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") @safetestset "ZygotePullback returns a NetworkParameters" include("pullbacks/zygote_pullback.jl") @safetestset "AutoEncoderLoss and a loss's NetworkParameters" include("loss/losses.jl") + @safetestset "apply_parametric and the concatenation of the system parameters" include("parametric.jl") + @safetestset "ParametricLoss and its ZygotePullback" include("loss/parametric_loss.jl") @safetestset "map_to_cpu" include("map_to_cpu.jl") @safetestset "changebackend" include("integration/changebackend_tests.jl") @safetestset "HDF5 save/load for GML special array types" include("integration/hdf5_support.jl") @@ -87,6 +89,7 @@ if "core" in GROUPS @safetestset "Classifier accuracy" include("data_loader/accuracy.jl") @safetestset "DataLoader batching" include("data_loader/batch.jl") @safetestset "DataLoader constructors" include("data_loader/data_loader.jl") + @safetestset "ParametricDataLoader batching and training" include("data_loader/parametric_data_loader.jl") @safetestset "Data loader docstring examples" include("integration/docstrings/data_loader.jl") @safetestset "Layer and architecture docstring examples" include("integration/docstrings/layers_and_architectures.jl") @safetestset "Loss docstring examples" include("integration/docstrings/losses.jl")