diff --git a/Project.toml b/Project.toml index e221589f4..0f0d53a0d 100644 --- a/Project.toml +++ b/Project.toml @@ -3,12 +3,6 @@ uuid = "194d25b2-d3f5-49f0-af24-c124f4aa80cc" version = "0.4.8" authors = ["Michael Kraus "] -[weakdeps] -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" - -[extensions] -HDF5Ext = "HDF5" - [deps] AbstractNeuralNetworks = "60874f82-5ada-4c70-bd1c-fa6be7711c8a" BandedMatrices = "aae01518-5342-5314-be14-df237901396f" @@ -19,6 +13,7 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GeometricBase = "9a0b12b7-583b-4f04-aa1f-d8551b6addc9" GeometricEquations = "c85262ba-a08a-430a-b926-d29770767bf2" +GeometricOptimizers = "fc236c15-5557-4942-aa65-b650f329279e" GeometricSolutions = "7843afe4-64f4-4df4-9231-049495c56661" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -37,6 +32,15 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" +[weakdeps] +HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" + +[sources] +GeometricOptimizers = {path = "../GeometricOptimizers"} + +[extensions] +HDF5Ext = "HDF5" + [compat] AbstractNeuralNetworks = "0.5, 0.6" BandedMatrices = "1" @@ -64,13 +68,13 @@ julia = "1.9" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" GeometricEquations = "c85262ba-a08a-430a-b926-d29770767bf2" GeometricIntegrators = "dcce2d33-59f6-5b8d-9047-0defad88ae06" GeometricProblems = "18cb22b4-ad41-5c80-9c5f-710df63fbdc9" +HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RungeKutta = "fb486d5c-30a0-4a8a-8415-a8b4ace5a6f7" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" diff --git a/docs/src/assets/custom.sty b/docs/src/assets/custom.sty index ba1673bd1..c2df8a857 100644 --- a/docs/src/assets/custom.sty +++ b/docs/src/assets/custom.sty @@ -125,7 +125,7 @@ % custom hyphenation \usepackage[htt]{hyphenat} -% renew texttt for better line breaks; taken from https://tex.stackexchange.com/questions/579789/combining-wrapping-texttt-with-sections-and-toc-improper-alphabetic-constan +% renew texttt for better line breaks; taken from https://tex.stackexchange.com/questions/579789/combining-wrapping-texttt-with-sections-and-toc-improper-alphabetic-constant \DeclareRobustCommand{\texttt}[1]{% \begingroup \ttfamily diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index c1761dcb0..435b92115 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -1,415 +1,404 @@ module GeometricMachineLearning - using AbstractNeuralNetworks - using BandedMatrices - using ChainRulesCore - using Distances - using GeometricBase - using GeometricSolutions: GeometricSolution, EnsembleSolution, DataSeries, StateVariable, TimeSeries - using GeometricEquations: EnsembleProblem, ODEProblem, HODEProblem, ODEEnsemble, HODEEnsemble - using KernelAbstractions - using LinearAlgebra - using NNlib - using ProgressMeter - using Random - using Zygote - using ForwardDiff - using InteractiveUtils - using TimerOutputs - import LazyArrays - import SymbolicNeuralNetworks - import SymbolicNeuralNetworks: input_dimension, output_dimension, SymbolicPullback - using SymbolicNeuralNetworks: derivative, _get_contents, _get_params, SymbolicNeuralNetwork - using Symbolics: @variables, substitute - - import AbstractNeuralNetworks: Architecture, Model, AbstractExplicitLayer, AbstractExplicitCell, AbstractNeuralNetwork , NeuralNetwork, UnknownArchitecture, FeedForwardLoss - import AbstractNeuralNetworks: Chain, GridCell - import AbstractNeuralNetworks: Dense, Linear, Recurrent - import AbstractNeuralNetworks: IdentityActivation, ZeroVector - import AbstractNeuralNetworks: add!, update! - import AbstractNeuralNetworks: layer - import AbstractNeuralNetworks: initialparameters - import AbstractNeuralNetworks: parameterlength - import AbstractNeuralNetworks: GlorotUniform - import AbstractNeuralNetworks: params, architecture, model, dim - import AbstractNeuralNetworks: AbstractPullback, NetworkLoss, _compute_loss - import AbstractNeuralNetworks: networkbackend - import AbstractNeuralNetworks: save, load - # export params, architetcure, model - export dim - import NNlib: σ, sigmoid, softmax - import Base: iterate, eltype - #import LogExpFunctions: softmax - - export CPU, GPU - export Chain, NeuralNetwork - export Dense, Linear - export initialparameters - export parameterlength - export NeuralNetworkParameters - - export σ, sigmoid, softmax - - # from GeometricBase to print docs - export description - - include("utils.jl") - - include("data_loader/data_loader.jl") - - # INCLUDE ARRAYS - include("arrays/skew_symmetric.jl") - include("arrays/symmetric.jl") - include("arrays/poisson_tensor.jl") - include("arrays/abstract_lie_algebra_horizontal.jl") - include("arrays/stiefel_lie_algebra_horizontal.jl") - include("arrays/grassmann_lie_algebra_horizontal.jl") - include("arrays/triangular.jl") - include("arrays/lower_triangular.jl") - include("arrays/upper_triangular.jl") - - export SymmetricMatrix, PoissonTensor, SkewSymMatrix - export StiefelLieAlgHorMatrix - export SymplecticLieAlgMatrix, SymplecticLieAlgHorMatrix - export GrassmannLieAlgHorMatrix - export StiefelProjection, SymplecticProjection - export LowerTriangular, UpperTriangular - - include("kernels/assign_q_and_p.jl") - include("kernels/tensor_mat_mul.jl") - include("kernels/tensor_tensor_mul.jl") - include("kernels/tensor_transpose_tensor_mul.jl") - include("kernels/tensor_tensor_transpose_mul.jl") - include("kernels/tensor_transpose_mat_mul.jl") - include("kernels/tensor_transpose_tensor_transpose_mul.jl") - include("kernels/mat_tensor_mul.jl") - include("kernels/tensor_transpose.jl") - include("kernels/exponentials/tensor_exponential.jl") - include("kernels/inverses/cpu_inverse.jl") - include("kernels/inverses/inverse_2x2.jl") - include("kernels/inverses/inverse_3x3.jl") - include("kernels/inverses/inverse_4x4.jl") - include("kernels/inverses/inverse_5x5.jl") - include("kernels/inverses/tensor_cayley.jl") - include("kernels/inverses/tensor_mat_skew_sym_assign.jl") - include("kernels/vec_tensor_mul.jl") - - include("kernels/kernel_ad_routines/assign_q_and_p.jl") - include("kernels/kernel_ad_routines/tensor_mat_mul.jl") - include("kernels/kernel_ad_routines/mat_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl") - include("kernels/kernel_ad_routines/tensor_transpose.jl") - include("kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl") - include("kernels/kernel_ad_routines/vec_tensor_mul.jl") - # export tensor_mat_mul - - # this defines empty retraction type structs (doesn't rely on anything) - include("optimizers/manifold_related/retraction_types.jl") - - export MatrixSoftmax, VectorSoftmax - include("activations/softmax.jl") - - # are these needed? - export UnknownProblem, NothingFunction - - # + operation has been overloaded to work with NamedTuples! - export _add, apply_toNT, split_and_flatten, add! - - # GPU specific operations - export convert_to_dev, Device, CPUDevice - - export GradientLayerQ, GradientLayerP, ActivationLayerQ, ActivationLayerP, LinearLayerQ, LinearLayerP - export Linear - export ResidualLayer - export LinearSymplecticLayerP, LinearSymplecticLayerQ - export SymplecticStiefelLayer - - include("manifolds/abstract_manifold.jl") - include("manifolds/stiefel_manifold.jl") - # include("manifolds/symplectic_stiefel_manifold.jl") - include("manifolds/grassmann_manifold.jl") - - include("arrays/stiefel_projection.jl") - - export StiefelManifold, SymplecticStiefelManifold, GrassmannManifold, Manifold - export rgrad, metric - - include("optimizers/optimizer_method.jl") - include("optimizers/optimizer_caches.jl") - include("optimizers/optimizer.jl") - include("optimizers/gradient_optimizer.jl") - include("optimizers/momentum_optimizer.jl") - include("optimizers/adam_optimizer.jl") - include("optimizers/adam_optimizer_with_learning_rate_decay.jl") - include("optimizers/bfgs_cache.jl") - include("optimizers/bfgs_optimizer.jl") - include("optimizers/init_optimizer_cache.jl") - - include("optimizers/manifold_related/global_sections.jl") - include("optimizers/manifold_related/modified_exponential.jl") - include("optimizers/manifold_related/retractions.jl") - - include("layers/sympnets.jl") - include("layers/bias_layer.jl") - include("layers/resnet.jl") - include("layers/manifold_layer.jl") - include("layers/stiefel_layer.jl") - include("layers/grassmann_layer.jl") - include("layers/multi_head_attention.jl") - include("layers/volume_preserving_attention.jl") - include("layers/volume_preserving_feedforward.jl") - include("layers/transformer.jl") - include("layers/psd_like_layer.jl") - include("layers/classification.jl") - - # include("layers/symplectic_stiefel_layer.jl") - export StiefelLayer, GrassmannLayer, ManifoldLayer - export PSDLayer - export MultiHeadAttention - export VolumePreservingAttention - export VolumePreservingFeedForwardLayer - export VolumePreservingLowerLayer - export VolumePreservingUpperLayer - export VolumePreservingTransformer - export NeuralNetworkIntegrator - export ResNet - export Transformer - export TransformerIntegrator, StandardTransformerIntegrator - - # INCLUDE OPTIMIZERS - export OptimizerMethod, AbstractCache - export GradientOptimizer, GradientCache - export MomentumOptimizer, MomentumCache - export AdamOptimizerWithDecay - export AdamOptimizer, AdamCache - export AdamOptimizerWithDecay - export BFGSOptimizer, BFGSCache - - export Optimizer - export optimization_step! - - export GlobalSection, apply_section, apply_section! - export global_rep - export Geodesic, Cayley - export geodesic, cayley - export retraction - # export ⊙², √ᵉˡᵉ, /ᵉˡᵉ, scalar_add - export update! - export check - - #INCLUDE ABSTRACT TRAINING integrator - export AbstractTrainingMethod - - export loss_single #, loss - - export HnnTrainingMethod - export LnnTrainingMethod - export SympNetTrainingMethod - - include("training_method/abstract_training_method.jl") - - # INCLUDE DATA TRAINING STRUCTURE - export AbstractDataShape, TrajectoryData, SampledData - export get_length_trajectory, get_Δt, get_nb_point, get_nb_trajectory, get_data - - include("data/data_shape.jl") - - export AbstractDataSymbol - export PositionSymbol, PhaseSpaceSymbol, DerivativePhaseSpaceSymbol, PosVeloAccSymbol, PosVeloSymbol - export DataSymbol - export can_reduce, symbols, symboldiff - - include("data/data_symbol.jl") - - # INCLUDE TRAINING INTEGRATOR - - export TrainingMethod - export symbol, shape - export min_length_batch - - - include("training_method/training_method.jl") - - # INCLUDE DATA TRAINING STRUCTURE - export AbstractTrainingData - export TrainingData - export shape, symbols, dim, noisemaker, data_symbols # , problem - export reduce_symbols, reshape_intoSampledData - export aresame - - include("data/data_training.jl") - - export get_batch, complete_batch_size, check_batch_size - - include("data/batch.jl") +using AbstractNeuralNetworks +using BandedMatrices +using ChainRulesCore +using Distances +using GeometricBase +using GeometricSolutions: GeometricSolution, EnsembleSolution, DataSeries, StateVariable, + TimeSeries +using GeometricEquations: EnsembleProblem, ODEProblem, HODEProblem, ODEEnsemble, + HODEEnsemble +using KernelAbstractions +using LinearAlgebra +using NNlib +using ProgressMeter +using Random +using Zygote +using ForwardDiff +using InteractiveUtils +using TimerOutputs +import LazyArrays +import SymbolicNeuralNetworks +import SymbolicNeuralNetworks: input_dimension, output_dimension, SymbolicPullback +using SymbolicNeuralNetworks: derivative, _get_contents, _get_params, SymbolicNeuralNetwork +using Symbolics: @variables, substitute + +using GeometricOptimizers +using GeometricOptimizers: OptimizerSolution, Cayley, Geodesic, cayley, geodesic, retraction, + apply_section, apply_section!, OptimizerMethod, + GradientCache, MomentumCache, AdamCache + +import AbstractNeuralNetworks: Architecture, Model, AbstractExplicitLayer, + AbstractExplicitCell, AbstractNeuralNetwork, NeuralNetwork, + UnknownArchitecture, FeedForwardLoss +import AbstractNeuralNetworks: Chain, GridCell +import AbstractNeuralNetworks: Dense, Linear, Recurrent +import AbstractNeuralNetworks: IdentityActivation, ZeroVector +import AbstractNeuralNetworks: add!, update! +import AbstractNeuralNetworks: layer +import AbstractNeuralNetworks: initialparameters +import AbstractNeuralNetworks: parameterlength +import AbstractNeuralNetworks: GlorotUniform +import AbstractNeuralNetworks: params, architecture, model, dim +import AbstractNeuralNetworks: AbstractPullback, NetworkLoss, _compute_loss +import AbstractNeuralNetworks: networkbackend +import AbstractNeuralNetworks: save, load +# export params, architetcure, model +export dim +import NNlib: σ, sigmoid, softmax +import Base: iterate, eltype +#import LogExpFunctions: softmax + +export CPU, GPU +export Chain, NeuralNetwork +export Dense, Linear +export initialparameters +export parameterlength +export NeuralNetworkParameters + +export σ, sigmoid, softmax + +# from GeometricBase to print docs +export description + +include("utils.jl") + +include("data_loader/data_loader.jl") + +# INCLUDE ARRAYS +include("arrays/skew_symmetric.jl") +include("arrays/symmetric.jl") +include("arrays/poisson_tensor.jl") +include("arrays/abstract_lie_algebra_horizontal.jl") +include("arrays/stiefel_lie_algebra_horizontal.jl") +include("arrays/grassmann_lie_algebra_horizontal.jl") +include("arrays/triangular.jl") +include("arrays/lower_triangular.jl") +include("arrays/upper_triangular.jl") + +export SymmetricMatrix, PoissonTensor, SkewSymMatrix +export StiefelLieAlgHorMatrix +export SymplecticLieAlgMatrix, SymplecticLieAlgHorMatrix +export GrassmannLieAlgHorMatrix +export StiefelProjection, SymplecticProjection +export LowerTriangular, UpperTriangular + +include("kernels/assign_q_and_p.jl") +include("kernels/tensor_mat_mul.jl") +include("kernels/tensor_tensor_mul.jl") +include("kernels/tensor_transpose_tensor_mul.jl") +include("kernels/tensor_tensor_transpose_mul.jl") +include("kernels/tensor_transpose_mat_mul.jl") +include("kernels/tensor_transpose_tensor_transpose_mul.jl") +include("kernels/mat_tensor_mul.jl") +include("kernels/tensor_transpose.jl") +include("kernels/exponentials/tensor_exponential.jl") +include("kernels/inverses/cpu_inverse.jl") +include("kernels/inverses/inverse_2x2.jl") +include("kernels/inverses/inverse_3x3.jl") +include("kernels/inverses/inverse_4x4.jl") +include("kernels/inverses/inverse_5x5.jl") +include("kernels/inverses/tensor_cayley.jl") +include("kernels/inverses/tensor_mat_skew_sym_assign.jl") +include("kernels/vec_tensor_mul.jl") + +include("kernels/kernel_ad_routines/assign_q_and_p.jl") +include("kernels/kernel_ad_routines/tensor_mat_mul.jl") +include("kernels/kernel_ad_routines/mat_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl") +include("kernels/kernel_ad_routines/tensor_transpose.jl") +include("kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl") +include("kernels/kernel_ad_routines/vec_tensor_mul.jl") +# export tensor_mat_mul + +export MatrixSoftmax, VectorSoftmax +include("activations/softmax.jl") + +# are these needed? +export UnknownProblem, NothingFunction + +# + operation has been overloaded to work with NamedTuples! +export _add, apply_toNT, split_and_flatten, add! + +# GPU specific operations +export convert_to_dev, Device, CPUDevice + +export GradientLayerQ, GradientLayerP, ActivationLayerQ, ActivationLayerP, LinearLayerQ, + LinearLayerP +export Linear +export ResidualLayer +export LinearSymplecticLayerP, LinearSymplecticLayerQ +export SymplecticStiefelLayer + +include("manifolds/abstract_manifold.jl") +include("manifolds/stiefel_manifold.jl") +# include("manifolds/symplectic_stiefel_manifold.jl") +include("manifolds/grassmann_manifold.jl") + +include("arrays/stiefel_projection.jl") + +export StiefelManifold, GrassmannManifold, Manifold +export rgrad, metric + +include("layers/sympnets.jl") +include("layers/bias_layer.jl") +include("layers/resnet.jl") +include("layers/manifold_layer.jl") +include("layers/stiefel_layer.jl") +include("layers/grassmann_layer.jl") +include("layers/multi_head_attention.jl") +include("layers/volume_preserving_attention.jl") +include("layers/volume_preserving_feedforward.jl") +include("layers/transformer.jl") +include("layers/psd_like_layer.jl") +include("layers/classification.jl") + +# include("layers/symplectic_stiefel_layer.jl") +export StiefelLayer, GrassmannLayer, ManifoldLayer +export PSDLayer +export MultiHeadAttention +export VolumePreservingAttention +export VolumePreservingFeedForwardLayer +export VolumePreservingLowerLayer +export VolumePreservingUpperLayer +export VolumePreservingTransformer +export NeuralNetworkIntegrator +export ResNet +export Transformer +export TransformerIntegrator, StandardTransformerIntegrator + +# INCLUDE OPTIMIZERS — types come from GeometricOptimizers +export OptimizerMethod, AbstractCache +export GradientMethod, GradientCache, GradientState +export MomentumMethod, MomentumCache, MomentumState +export Adam, AdamCache, AdamState +export Optimizer +export optimization_step! +export GlobalSection, apply_section, apply_section! +export global_rep +export Geodesic, Cayley +export geodesic, cayley +export retraction +export update! +export check +# backward-compat aliases (old names → new names) +const GradientOptimizer = GradientMethod +const MomentumOptimizer = MomentumMethod +const AdamOptimizer = Adam +export GradientOptimizer, MomentumOptimizer, AdamOptimizer +export AdamOptimizerWithDecay + +#INCLUDE ABSTRACT TRAINING integrator +export AbstractTrainingMethod + +export loss_single #, loss + +export HnnTrainingMethod +export LnnTrainingMethod +export SympNetTrainingMethod + +include("training_method/abstract_training_method.jl") + +# INCLUDE DATA TRAINING STRUCTURE +export AbstractDataShape, TrajectoryData, SampledData +export get_length_trajectory, get_Δt, get_nb_point, get_nb_trajectory, get_data + +include("data/data_shape.jl") + +export AbstractDataSymbol +export PositionSymbol, PhaseSpaceSymbol, DerivativePhaseSpaceSymbol, PosVeloAccSymbol, + PosVeloSymbol +export DataSymbol +export can_reduce, symbols, symboldiff + +include("data/data_symbol.jl") + +# INCLUDE TRAINING INTEGRATOR + +export TrainingMethod +export symbol, shape +export min_length_batch + +include("training_method/training_method.jl") + +# INCLUDE DATA TRAINING STRUCTURE +export AbstractTrainingData +export TrainingData +export shape, symbols, dim, noisemaker, data_symbols # , problem +export reduce_symbols, reshape_intoSampledData +export aresame + +include("data/data_training.jl") + +export get_batch, complete_batch_size, check_batch_size + +include("data/batch.jl") + +# INCLUDE BACKENDS +export LuxBackend +export NeuralNetwork +export arch + +include("backends/backends.jl") +include("backends/lux.jl") + +export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss + +#INCLUDE ARCHITECTURES +include("architectures/neural_network_integrator.jl") +include("architectures/resnet.jl") +include("architectures/transformer_integrator.jl") +include("architectures/standard_transformer_integrator.jl") +include("architectures/sympnet.jl") +include("architectures/autoencoder.jl") +include("architectures/symplectic_autoencoder.jl") +include("architectures/psd.jl") +include("architectures/fixed_width_network.jl") +include("architectures/hamiltonian_neural_network.jl") +include("architectures/lagrangian_neural_network.jl") +include("architectures/variable_width_network.jl") +include("architectures/recurrent_neural_network.jl") +include("architectures/LSTM_neural_network.jl") +include("architectures/transformer_neural_network.jl") +include("architectures/volume_preserving_feedforward.jl") +include("architectures/volume_preserving_transformer.jl") - # INCLUDE BACKENDS - export LuxBackend - export NeuralNetwork - export arch +export HamiltonianArchitecture +export LagrangianNeuralNetwork +export SympNet, LASympNet, GSympNet +export RecurrentNeuralNetwork +export LSTMNeuralNetwork +export ClassificationTransformer, ClassificationLayer +export VolumePreservingFeedForward +export SymplecticAutoencoder, PSDArch +export HamiltonianArchitecture, StandardHamiltonianArchitecture, + GeneralizedHamiltonianArchitecture - include("backends/backends.jl") - include("backends/lux.jl") +export solve!, encoder, decoder - export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss +export train!, apply!, jacobian! +export iterate - #INCLUDE ARCHITECTURES - include("architectures/neural_network_integrator.jl") - include("architectures/resnet.jl") - include("architectures/transformer_integrator.jl") - include("architectures/standard_transformer_integrator.jl") - include("architectures/sympnet.jl") - include("architectures/autoencoder.jl") - include("architectures/symplectic_autoencoder.jl") - include("architectures/psd.jl") - include("architectures/fixed_width_network.jl") - include("architectures/hamiltonian_neural_network.jl") - include("architectures/lagrangian_neural_network.jl") - include("architectures/variable_width_network.jl") - include("architectures/recurrent_neural_network.jl") - include("architectures/LSTM_neural_network.jl") - include("architectures/transformer_neural_network.jl") - include("architectures/volume_preserving_feedforward.jl") - include("architectures/volume_preserving_transformer.jl") +export default_arch - export HamiltonianArchitecture - export LagrangianNeuralNetwork - export SympNet, LASympNet, GSympNet - export RecurrentNeuralNetwork - export LSTMNeuralNetwork - export ClassificationTransformer, ClassificationLayer - export VolumePreservingFeedForward - export SymplecticAutoencoder, PSDArch - export HamiltonianArchitecture, StandardHamiltonianArchitecture, GeneralizedHamiltonianArchitecture +include("architectures/default_architecture.jl") - export solve!, encoder, decoder +include("loss/losses.jl") +include("loss/hnn_loss.jl") - export train!, apply!, jacobian! - export iterate +export AbstractPullback, ZygotePullback, SymbolicPullback +include("pullbacks/zygote_pullback.jl") +include("pullbacks/symbolic_hnn_pullback.jl") - export default_arch +export DataLoader, onehotbatch +export Batch, optimize_for_one_epoch! +include("data_loader/tensor_assign.jl") +include("data_loader/matrix_assign.jl") +include("data_loader/mnist_utils.jl") +include("data_loader/batch.jl") +include("data_loader/optimize.jl") - include("architectures/default_architecture.jl") +# INCLUDE TRAINING parameters - include("loss/losses.jl") - include("loss/hnn_loss.jl") +export TrainingParameters - export AbstractPullback, ZygotePullback, SymbolicPullback - include("pullbacks/zygote_pullback.jl") - include("pullbacks/symbolic_hnn_pullback.jl") +include("training/training_parameters.jl") - export DataLoader, onehotbatch - export Batch, optimize_for_one_epoch! - include("data_loader/tensor_assign.jl") - include("data_loader/matrix_assign.jl") - include("data_loader/mnist_utils.jl") - include("data_loader/batch.jl") - include("data_loader/optimize.jl") +# INCLUDE NEURALNET SOLUTION - export default_optimizer +export SingleHistory +export parameters, datashape +export History +export last, sizemax, nbtraining, show - include("optimizers/default_optimizer.jl") +include("nnsolution/history.jl") - # INCLUDE TRAINING parameters +export NeuralNetSolution +export problem, timestep, history, size_history +export set_sizemax_history - export TrainingParameters +include("nnsolution/neural_net_solution.jl") - include("training/training_parameters.jl") +export EnsembleNeuralNetSolution +export push!, merge! - # INCLUDE NEURALNET SOLUTION +include("nnsolution/neural_net_solution_ensemble.jl") - export SingleHistory - export parameters, datashape - export History - export last, sizemax, nbtraining, show +# INCLUDE TRAINING integrator - include("nnsolution/history.jl") +export TrainingSet +export parameters # , data - export NeuralNetSolution - export problem, timestep, history, size_history - export set_sizemax_history +include("training/training_set.jl") - include("nnsolution/neural_net_solution.jl") +export EnsembleTraining +export isnnShared, isParametersShared, isDataShared +export parameters, data +export push!, merge!, size - export EnsembleNeuralNetSolution - export push!, merge! +include("training/ensemble_training.jl") - include("nnsolution/neural_net_solution_ensemble.jl") +include("training/nn_parameters_transformation.jl") - # INCLUDE TRAINING integrator +export loss_gradient +export train! - export TrainingSet - export parameters # , data +include("training/train.jl") - include("training/training_set.jl") +export SymplecticEuler +export SymplecticEulerA, SymplecticEulerB +export SEuler, SEulerA, SEulerB - export EnsembleTraining - export isnnShared, isParametersShared, isDataShared - export parameters, data - export push!, merge!, size +include("training_method/symplectic_euler.jl") - include("training/ensemble_training.jl") +export HnnExactMethod +export ExactHnn - include("training/nn_parameters_transformation.jl") +include("training_method/hnn_exact_method.jl") - export loss_gradient - export train! +export VariationalMethod +export VariationalMidPointMethod +export VariaMidPoint - include("training/train.jl") +include("training_method/variational_method.jl") - export SymplecticEuler - export SymplecticEulerA, SymplecticEulerB - export SEuler, SEulerA, SEulerB +export LnnExactMethod +export ExactLnn - include("training_method/symplectic_euler.jl") +include("training_method/lnn_exact_method.jl") - export HnnExactMethod - export ExactHnn +export BasicSympNetMethod +export BasicSympNet - include("training_method/hnn_exact_method.jl") +include("training_method/sympnet_basic_method.jl") - export VariationalMethod - export VariationalMidPointMethod - export VariaMidPoint +export default_method - include("training_method/variational_method.jl") +include("training/default_method.jl") - export LnnExactMethod - export ExactLnn +# INCLUDE ASSERTION Function +export matching +include("training/matching.jl") - include("training_method/lnn_exact_method.jl") +include("reduced_system/reduced_system.jl") - export BasicSympNetMethod - export BasicSympNet +export HRedSys, reduction_error, projection_error, integrate_reduced_system, + integrate_full_system - include("training_method/sympnet_basic_method.jl") +include("layers/linear_symplectic_attention.jl") +include("layers/symplectic_attention.jl") +include("architectures/linear_symplectic_transformer.jl") +include("architectures/symplectic_transformer.jl") - export default_method +export LinearSymplecticAttention, LinearSymplecticAttentionQ, LinearSymplecticAttentionP +export LinearSymplecticTransformer +export SymplecticAttention, SymplecticAttentionQ, SymplecticAttentionP +export SymplecticTransformer - include("training/default_method.jl") +include("map_to_cpu.jl") - - # INCLUDE ASSERTION Function - export matching - include("training/matching.jl") - - include("reduced_system/reduced_system.jl") - - export HRedSys, reduction_error, projection_error, integrate_reduced_system, integrate_full_system - - include("layers/linear_symplectic_attention.jl") - include("layers/symplectic_attention.jl") - include("architectures/linear_symplectic_transformer.jl") - include("architectures/symplectic_transformer.jl") - - export LinearSymplecticAttention, LinearSymplecticAttentionQ, LinearSymplecticAttentionP - export LinearSymplecticTransformer - export SymplecticAttention, SymplecticAttentionQ, SymplecticAttentionP - export SymplecticTransformer - - include("map_to_cpu.jl") - - export save, load +export save, load end diff --git a/src/arrays/grassmann_lie_algebra_horizontal.jl b/src/arrays/grassmann_lie_algebra_horizontal.jl index cb7414c63..0e5e455bb 100644 --- a/src/arrays/grassmann_lie_algebra_horizontal.jl +++ b/src/arrays/grassmann_lie_algebra_horizontal.jl @@ -197,4 +197,9 @@ function _round(B::GrassmannLieAlgHorMatrix; kwargs...) B.N, B.n ) +end + +function Base.copyto!(A::GrassmannLieAlgHorMatrix, B::GrassmannLieAlgHorMatrix) + copyto!(A.B, B.B) + A end \ No newline at end of file diff --git a/src/arrays/skew_symmetric.jl b/src/arrays/skew_symmetric.jl index f3ae66e4d..440928bcb 100644 --- a/src/arrays/skew_symmetric.jl +++ b/src/arrays/skew_symmetric.jl @@ -145,8 +145,8 @@ end Base.:*(α::Real, A::SkewSymMatrix) = A*α -function Base.zeros(ST::Type{SkewSymMatrix{<:Real}}, n::Int) - zeros(CPU(), ST, n) +function Base.zeros(::Type{SkewSymMatrix{T}}, n::Int) where T + zeros(CPU(), SkewSymMatrix{T}, n) end function Base.zeros(backend::KernelAbstractions.Backend, ::Type{SkewSymMatrix{T}}, n::Int) where T @@ -342,7 +342,9 @@ function _round(A::AbstractArray; kwargs...) round.(A; kwargs...) end -# define routines for generalizing ChainRulesCore to SkewSymMatrix +Base.fill!(A::SkewSymMatrix, val) = (fill!(A.S, val); A) + +# define routines for generalizing ChainRulesCore to SkewSymMatrix ChainRulesCore.ProjectTo(A::SkewSymMatrix) = ProjectTo{SkewSymMatrix}(; skew_sym = ProjectTo(A.S)) (project::ProjectTo{SkewSymMatrix})(dA::AbstractMatrix) = SkewSymMatrix(project.skew_sym(map_to_Skew(dA)), size(dA, 2)) (project::ProjectTo{SkewSymMatrix})(dA::SkewSymMatrix) = SkewSymMatrix(project.skew_sym(dA.S), dA.n) \ No newline at end of file diff --git a/src/arrays/stiefel_lie_algebra_horizontal.jl b/src/arrays/stiefel_lie_algebra_horizontal.jl index dd144b09f..d77e552df 100644 --- a/src/arrays/stiefel_lie_algebra_horizontal.jl +++ b/src/arrays/stiefel_lie_algebra_horizontal.jl @@ -317,4 +317,6 @@ function _round(B::StiefelLieAlgHorMatrix; kwargs...) B.N, B.n ) -end \ No newline at end of file +end + +Base.fill!(A::StiefelLieAlgHorMatrix, val) = (fill!(A.A, val); fill!(A.B, val); A) \ No newline at end of file diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index a575ab31b..c10fbacae 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -30,7 +30,7 @@ batch(dl) # output [ Info: You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) number of parameters. -([(1, 5), (1, 3)], [(1, 4), (1, 1)], [(1, 2)]) +([(1, 3), (1, 1)], [(1, 5), (1, 4)], [(1, 2)]) ``` Here the first index is always 1 (the time dimension). We get a total number of 3 batches. @@ -125,8 +125,8 @@ println(stdout, batch(dl₁), "\n", batch(dl₂)) Number of batches of dl₁: 2 Number of batches of dl₂: 2 -([(1, 1), (4, 1), (2, 1)], [(3, 1)]) -([(1, 3), (1, 2), (1, 4)], [(1, 1), (1, 5)]) +([(1, 1), (3, 1), (2, 1)], [(4, 1)]) +([(1, 5), (1, 4), (1, 1)], [(1, 3), (1, 2)]) ``` Here we see that in the *autoencoder case* that last minibatch has an additional element. diff --git a/src/loss/hnn_loss.jl b/src/loss/hnn_loss.jl index 3ffd800e5..397e0602d 100644 --- a/src/loss/hnn_loss.jl +++ b/src/loss/hnn_loss.jl @@ -26,15 +26,17 @@ function HNNLoss(arch::HamiltonianArchitecture) HNNLoss(hamiltonian_vector_field(arch)) end -function (loss::HNNLoss)( ::Union{Chain, AbstractExplicitLayer}, - ps::Union{NeuralNetworkParameters, NamedTuple}, - input::QPTOAT, - output::QPTOAT) +AbstractNeuralNetworks.NetworkLoss(arch::HamiltonianArchitecture) = HNNLoss(arch) + +function (loss::HNNLoss)(::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::QPTOAT, + output::QPTOAT) loss(ps, input, output) end -function (loss::HNNLoss)( ps::Union{NeuralNetworkParameters, NamedTuple}, - input::QPTOAT, - output::QPTOAT) +function (loss::HNNLoss)(ps::Union{NeuralNetworkParameters, NamedTuple}, + input::QPTOAT, + output::QPTOAT) norm(loss.hvf(input, ps) - output) / norm(output) -end \ No newline at end of file +end diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 93764f99a..375787e75 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -1,7 +1,10 @@ +# type piracy! This should go into `AbstractNeuralNetworks`! ... +AbstractNeuralNetworks.NetworkLoss(nn::NeuralNetwork) = NetworkLoss(architecture(nn)) + @doc raw""" TransformerLoss(seq_length, prediction_window) -Make an instance of the transformer loss. +Make an instance of the transformer loss. This should be used together with a neural network of type [`TransformerIntegrator`](@ref). @@ -9,7 +12,7 @@ This should be used together with a neural network of type [`TransformerIntegrat `TransformerLoss` applies a neural network to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -40,7 +43,7 @@ So `TransformerLoss` simply does: ```math \mathtt{loss}(\mathcal{NN}, \mathtt{input}, \mathtt{output}) = || \mathcal{NN}(\mathtt{input})[(\mathtt{sl} - \mathtt{pw} + 1):\mathtt{end}] - \mathtt{output} || / || \mathtt{output} ||, ``` -where ``||\cdot||`` is the ``L_2`` norm. +where ``||\cdot||`` is the ``L_2`` norm. # Parameters @@ -52,14 +55,24 @@ struct TransformerLoss <: NetworkLoss prediction_window::Int end +function AbstractNeuralNetworks.NetworkLoss( + arch::TransformerIntegrator, prediction_window::Integer = 1) + TransformerLoss(arch.seq_length, prediction_window) +end + TransformerLoss(seq_length::Int) = TransformerLoss(seq_length, seq_length) -# This crops the output array of the neural network so that it conforms with the output it should be compared to. This is needed for the transformer loss. -function crop_array_for_transformer_loss(nn_output::AT, output::BT) where {T, T2, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T2, 3}} - @view nn_output[axes(output, 1), axes(output, 2) .+ size(nn_output, 2) .- size(output, 2), axes(output, 3)] +# This crops the output array of the neural network so that it conforms with the output it should be compared to. This is needed for the transformer loss. +function crop_array_for_transformer_loss(nn_output::AT, + output::BT) where {T, T2, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T2, 3}} + @view nn_output[ + axes(output, 1), axes(output, 2) .+ size(nn_output, 2) .- size(output, 2), + axes(output, 3)] end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, output::AT) where {T, AT <: AbstractArray{T, 3}} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, + output::AT) where {T, AT <: AbstractArray{T, 3}} input_dim, input_seq_length = size(input) output_dim, output_prediction_window = size(output) @assert input_dim == output_dim @@ -67,39 +80,46 @@ function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps: @assert output_prediction_window == loss.prediction_window predicted_output_uncropped = model(input, ps) - predicted_output_cropped = crop_array_for_transformer_loss(predicted_output_uncropped, output) + predicted_output_cropped = crop_array_for_transformer_loss( + predicted_output_uncropped, output) _compute_loss(predicted_output_cropped, output) end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, output::AT) where {T, AT <: AbstractArray{T, 2}} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::AT, + output::AT) where {T, AT <: AbstractArray{T, 2}} loss(model, ps, reshape(input, size(input)..., 1), reshape(output, size(output)..., 1)) end -function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::T, output::T) where {T <: QPT} +function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::T, output::T) where {T <: QPT} loss(model, ps, vcat(input.q, input.p), vcat(output.q, output.p)) end # @doc raw""" # ClassificationTransformerLoss() -# +# # Make an instance of `ClassificationTransformerLoss`. -# +# # This is to be used together with a [`ClassificationTransformer`](@ref). -# +# # It takes an input, parses it to the transformer and then crops it to conform with the desired output size. -# +# # Suppose the input is of dimension ``\mathtt{td}\times\mathtt{sl}``, where `td` is *transformer dimension* and `sl` is *sequence length*. -# The output of the transformer will again be of the same dimension: -# +# The output of the transformer will again be of the same dimension: +# # ```math # \mathrm{output}\in\mathbb{R}^{\mathtt{td}\times\mathtt{sl}}. # ``` -# +# # if the output dimension `cl` of the [`ClassificationLayer`](@ref) is differnt form `td`. # """ struct ClassificationTransformerLoss <: NetworkLoss end -function (loss::ClassificationTransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::AbstractArray, output::AbstractArray) +function (loss::ClassificationTransformerLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, + input::AbstractArray, output::AbstractArray) predicted_output_uncropped = model(input, ps) # predicted_output_cropped = crop_array_for_transformer_loss(predicted_output_uncropped, output) norm(predicted_output_uncropped - output) / norm(output) @@ -110,13 +130,13 @@ end Make an instance of `AutoEncoderLoss`. -This loss should always be used together with a neural network of type [`AutoEncoder`](@ref) (and it is also the default for training such a network). +This loss should always be used together with a neural network of type [`AutoEncoder`](@ref) (and it is also the default for training such a network). # Example `AutoEncoderLoss` applies a neural network to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -142,23 +162,29 @@ So `AutoEncoderLoss` simply does: ```math \mathtt{loss}(\mathcal{NN}, \mathtt{input}) = || \mathcal{NN}(\mathtt{input}) - \mathtt{input} || / || \mathtt{input} ||, ``` -where ``||\cdot||`` is the ``L_2`` norm. +where ``||\cdot||`` is the ``L_2`` norm. # Parameters This loss does not have any parameters. """ -struct AutoEncoderLoss <: NetworkLoss end +struct AutoEncoderLoss <: NetworkLoss end + +AbstractNeuralNetworks.NetworkLoss(::AutoEncoder) = AutoEncoderLoss() function (loss::AutoEncoderLoss)(nn::NeuralNetwork, input::QPTOAT) loss(nn.model, params(nn), input, input) end -function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT) +function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT) loss(model, ps, input, input) end -(loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT) = FeedForwardLoss()(model, ps, input, output) +function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer}, + ps::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT) + FeedForwardLoss()(model, ps, input, output) +end @doc raw""" ReducedLoss(encoder, decoder) @@ -171,7 +197,7 @@ This loss should be used together with a [`NeuralNetworkIntegrator`](@ref) or [` `ReducedLoss` applies the *encoder*, *integrator* and *decoder* neural networks in this order to an input and compares it to the `output` via an ``L_2`` norm: -```jldoctest +```jldoctest using GeometricMachineLearning using LinearAlgebra: norm import Random @@ -196,7 +222,7 @@ loss(transformer, input_mat, output_mat) ≈ norm(output_mat - output_prediction true ``` -So the loss computes: +So the loss computes: ```math \mathrm{loss}_{\mathcal{E}, \mathcal{D}}(\mathcal{NN}, \mathrm{input}, \mathrm{output}) = ||\mathcal{D}(\mathcal{NN}(\mathcal{E}(\mathrm{input}))) - \mathrm{output}||, @@ -204,7 +230,8 @@ So the loss computes: where ``\mathcal{E}`` is the [`Encoder`](@ref), ``\mathcal{D}`` is the [`Decoder`](@ref). ``\mathcal{NN}`` is the neural network we compute the loss of. """ -struct ReducedLoss{ET <: NeuralNetwork{<:Encoder}, DT <: NeuralNetwork{<:Decoder}} <: NetworkLoss +struct ReducedLoss{ET <: NeuralNetwork{<:Encoder}, DT <: NeuralNetwork{<:Decoder}} <: + NetworkLoss encoder::ET decoder::DT end @@ -224,6 +251,7 @@ function ReducedLoss(autoencoder::NeuralNetwork{<:AutoEncoder}) ReducedLoss(encoder(autoencoder), decoder(autoencoder)) end -function (loss::ReducedLoss)(model::Chain, params::NeuralNetworkParameters, input::CT, output::CT) where {CT <: QPTOAT} +function (loss::ReducedLoss)(model::Chain, params::NeuralNetworkParameters, + input::CT, output::CT) where {CT <: QPTOAT} _compute_loss(loss.decoder(model(loss.encoder(input), params)), output) -end \ No newline at end of file +end diff --git a/src/manifolds/grassmann_manifold.jl b/src/manifolds/grassmann_manifold.jl index 81e5e687d..0ee8972c5 100644 --- a/src/manifolds/grassmann_manifold.jl +++ b/src/manifolds/grassmann_manifold.jl @@ -117,4 +117,58 @@ function Ω(Y::GrassmannManifold{T}, Δ::AbstractMatrix{T}) where T # E = StiefelProjection(Y) # SkewSymMatrix(ΩSt - E * E' * ΩSt * E * E') SkewSymMatrix(ΩSt) +end + +function Base.copyto!(A::GrassmannManifold, B::GrassmannManifold) + A.A .= B.A + nothing +end + +Base.copy(A::GrassmannManifold) = GrassmannManifold(copy(A.A)) +Base.similar(A::GrassmannManifold) = GrassmannManifold(similar(A.A)) + +function GeometricOptimizers.global_rep( + λY::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Δ::AbstractMatrix{T} +) where T + N, n = size(λY.Y) + GrassmannLieAlgHorMatrix( + λY.λ' * Δ, + N, n + ) +end + +function GeometricOptimizers.update_section!( + Λᵗ::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Λ⁽ᵗ⁻¹⁾::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + B⁽ᵗ⁻¹⁾::AbstractMatrix{T}, + retraction +) where T + N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n + expB = retraction(B⁽ᵗ⁻¹⁾) + expB.A .= Λ⁽ᵗ⁻¹⁾.Y.A * expB.A[1:n, :] .+ Λ⁽ᵗ⁻¹⁾.λ * expB.A[(n+1):N, :] + Λᵗ.Y.A .= @view expB.A[:, 1:n] + Λᵗ.λ .= @view expB.A[:, (n+1):N] + nothing +end + +function GeometricOptimizers.cayley(B::GrassmannLieAlgHorMatrix{T}) where T + backend = networkbackend(B) + E = StiefelProjection(B) + 𝕆 = KernelAbstractions.zeros(backend, T, B.n, B.n) + 𝕀_small = one(𝕆) + 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) + 𝕀_big = one(B) + B̂ = hcat(vcat(𝕆, B.B), E) + B̄ = hcat(vcat(𝕀_small, 𝕆), vcat(zero(B.B'), -B.B'))' + GrassmannManifold((𝕀_big + T(0.5) * B̂ * inv(𝕀_small2 - T(0.5) * B̄' * B̂) * B̄') * (𝕀_big + T(0.5) * B)) +end + +function GeometricOptimizers._copyto!( + Λ₁::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}}, + Λ₂::GeometricOptimizers.GlobalSection{T, <:GrassmannManifold{T}} +) where T + copyto!(Λ₁.Y, Λ₂.Y) + copyto!(Λ₁.λ, Λ₂.λ) + Λ₁ end \ No newline at end of file diff --git a/src/manifolds/stiefel_manifold.jl b/src/manifolds/stiefel_manifold.jl index 5caf4afc0..bcd80e1ad 100644 --- a/src/manifolds/stiefel_manifold.jl +++ b/src/manifolds/stiefel_manifold.jl @@ -181,4 +181,88 @@ end function Base.copyto!(A::StiefelManifold, B::StiefelManifold) A.A .= B.A nothing +end + +Base.copy(A::StiefelManifold) = StiefelManifold(copy(A.A)) +Base.similar(A::StiefelManifold) = StiefelManifold(similar(A.A)) + +function Base.zero(Y::StiefelManifold{T}) where T + N, n = size(Y) + backend = KernelAbstractions.get_backend(Y.A) + zeros(backend, StiefelLieAlgHorMatrix{T}, N, n) +end + +# Bridge GML's StiefelManifold into GO's manifold optimization infrastructure. +# GO's methods dispatch on MT<:GO.Manifold{T}, but GML.StiefelManifold<:GML.Manifold{T} +# (a different type hierarchy), so we extend GO's functions explicitly for GML's type. + +function GeometricOptimizers.global_rep( + λY::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Δ::AbstractMatrix{T} +) where T + N, n = size(λY.Y) + StiefelLieAlgHorMatrix( + SkewSymMatrix(λY.Y.A' * Δ), + λY.λ' * Δ, + N, n + ) +end + +function GeometricOptimizers.update_section!( + Λᵗ::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Λ⁽ᵗ⁻¹⁾::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + B⁽ᵗ⁻¹⁾::AbstractMatrix{T}, + retraction +) where T + N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n + expB = retraction(B⁽ᵗ⁻¹⁾) + expB.A .= Λ⁽ᵗ⁻¹⁾.Y.A * expB.A[1:n, :] .+ Λ⁽ᵗ⁻¹⁾.λ * expB.A[(n+1):N, :] + Λᵗ.Y.A .= @view expB.A[:, 1:n] + Λᵗ.λ .= @view expB.A[:, (n+1):N] + nothing +end + +function GeometricOptimizers.cayley(B::StiefelLieAlgHorMatrix{T}) where T + E = StiefelProjection(B) + 𝕀_small = one(B.A) + 𝕆 = zero(𝕀_small) + 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) + 𝕀_big = one(B) + A_mat = B.A * 𝕀_small + B̂ = hcat(vcat(T(0.5) * A_mat, B.B), E) + B̄ = hcat(vcat(𝕀_small, T(0.5) * A_mat), vcat(zero(B.B'), -B.B'))' + StiefelManifold((𝕀_big + T(0.5) * B̂ * inv(𝕀_small2 - T(0.5) * B̄' * B̂) * B̄') * (𝕀_big + T(0.5) * B)) +end + +function GeometricOptimizers._copyto!( + Λ₁::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}}, + Λ₂::GeometricOptimizers.GlobalSection{T, <:StiefelManifold{T}} +) where T + copyto!(Λ₁.Y, Λ₂.Y) + copyto!(Λ₁.λ, Λ₂.λ) + Λ₁ +end + +# GO arithmetic bridges for GML's SkewSymMatrix (GO dispatches by module, not structure). +GeometricOptimizers._add!(a::SkewSymMatrix{T}, b::SkewSymMatrix{T}) where T = (a.S .+= b.S; a) +GeometricOptimizers._add!(a::SkewSymMatrix{T}, b::T) where T = (a.S .+= b; a) +GeometricOptimizers._rac!(B::SkewSymMatrix, A::SkewSymMatrix) = (B.S .= sqrt.(A.S); B) +GeometricOptimizers._square!(B::SkewSymMatrix, A::SkewSymMatrix) = (B.S .= A.S .^ 2; B) +GeometricOptimizers._div!(C::SkewSymMatrix, A::SkewSymMatrix, B::SkewSymMatrix) = (C.S .= A.S ./ B.S; C) + +# GO arithmetic bridges for GML's StiefelLieAlgHorMatrix. +function GeometricOptimizers._add!(A::StiefelLieAlgHorMatrix{T}, B::StiefelLieAlgHorMatrix{T}) where T + GeometricOptimizers._add!(A.A, B.A); A.B .+= B.B; A +end +function GeometricOptimizers._add!(A::StiefelLieAlgHorMatrix{T}, b::T) where T + GeometricOptimizers._add!(A.A, b); A.B .+= b; A +end +function GeometricOptimizers._rac!(B::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix) + GeometricOptimizers._rac!(B.A, A.A); B.B .= sqrt.(A.B); B +end +function GeometricOptimizers._square!(B::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix) + GeometricOptimizers._square!(B.A, A.A); B.B .= A.B .^ 2; B +end +function GeometricOptimizers._div!(C::StiefelLieAlgHorMatrix, A::StiefelLieAlgHorMatrix, B::StiefelLieAlgHorMatrix) + GeometricOptimizers._div!(C.A, A.A, B.A); C.B .= A.B ./ B.B; C end \ No newline at end of file diff --git a/src/optimizers/adam_optimizer.jl b/src/optimizers/adam_optimizer.jl deleted file mode 100644 index 27781769f..000000000 --- a/src/optimizers/adam_optimizer.jl +++ /dev/null @@ -1,56 +0,0 @@ -@doc raw""" - AdamOptimizer(η, ρ₁, ρ₂, δ) - -Make an instance of the Adam Optimizer. - -Here the cache consists of first and second moments that are updated as - -```math -B_1 \gets ((\rho_1 - \rho_1^t)/(1 - \rho_1^t))\cdot{}B_1 + (1 - \rho_1)/(1 - \rho_1^t)\cdot{}\nabla{}L, -``` -and - -```math -B_2 \gets ((\rho_2 - \rho_1^t)/(1 - \rho_2^t))\cdot{}B_2 + (1 - \rho_2)/(1 - \rho_2^t)\cdot\nabla{}L\odot\nabla{}L. -``` -The final velocity is computed as: - -```math -\mathrm{velocity} \gets -\eta{}B_1/\sqrt{B_2 + \delta}. -``` - -# Implementation - -The *velocity* is stored in the input to save memory: - -```julia -mul!(B, -o.method.η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -``` -where `B` is the input to the [`update!`] function. - -The algorithm and suggested defaults are taken from [goodfellow2016deep; page 301](@cite). -""" -struct AdamOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - ρ₁::T - ρ₂::T - δ::T - - AdamOptimizer(η = 1f-3, ρ₁ = 9f-1, ρ₂ = 9.9f-1, δ = 3f-7; T=typeof(η)) = new{T}(T(η), T(ρ₁), T(ρ₂), T(δ)) -end - -function AdamOptimizer(T::Type) - AdamOptimizer(T(1f-3)) -end - -function update!(o::Optimizer{<:AdamOptimizer{T}}, C::AdamCache, B::AbstractArray) where T - add!(C.B₁, ((o.method.ρ₁ - o.method.ρ₁^o.step)/(T(1.) - o.method.ρ₁^o.step))*C.B₁, ((T(1.) - o.method.ρ₁)/(T(1.) - o.method.ρ₁^o.step))*B) - add!(C.B₂, ((o.method.ρ₂ - o.method.ρ₂^o.step)/(T(1.) - o.method.ρ₂^o.step))*C.B₂, ((T(1.) - o.method.ρ₂)/(T(1.) - o.method.ρ₂^o.step))*⊙²(B)) - mul!(B, -o.method.η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -end - -# defaults: -⊙²(A::AbstractVecOrMat) = A.^2 -racᵉˡᵉ(A::AbstractVecOrMat) = sqrt.(A) -/ᵉˡᵉ(A::AbstractVecOrMat, B::AbstractVecOrMat) = A./B -scalar_add(A::AbstractVecOrMat, δ::Real) = A .+ δ \ No newline at end of file diff --git a/src/optimizers/adam_optimizer_with_learning_rate_decay.jl b/src/optimizers/adam_optimizer_with_learning_rate_decay.jl deleted file mode 100644 index f2faefea6..000000000 --- a/src/optimizers/adam_optimizer_with_learning_rate_decay.jl +++ /dev/null @@ -1,41 +0,0 @@ -@doc raw""" - AdamOptimizerWithDecay(n_epochs, η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8) - -Make an instance of the Adam Optimizer with weight decay. - -All except the first argument (the number of epochs) have defaults. - -The difference to the standard [`AdamOptimizer`](@ref) is that we change the learning reate ``\eta`` in each step. -Apart from the *time dependency* of ``\eta`` the two algorithms are however equivalent. -``\eta(0)`` starts with a high value ``\eta_1`` and then exponentially decrease until it reaches ``\eta_2`` with - -```math - \eta(t) = \gamma^t\eta_1, -``` -where ``\gamma = \exp(\log(\eta_1 / \eta_2) / \mathtt{n\_epochs}).`` -""" -struct AdamOptimizerWithDecay{T<:Real} <: OptimizerMethod{T} - η₁::T - η₂::T - ρ₁::T - ρ₂::T - δ::T - γ::T - n_epochs::Int - - function AdamOptimizerWithDecay(n_epochs::Int, η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8; T=typeof(η₁)) - γ = exp(log(η₂ / η₁) / n_epochs) - new{T}(T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ), T(γ), n_epochs) - end -end - -function AdamOptimizerWithDecay(n_epochs::Int, T::Type; η₁=1f-2, η₂=1f-6, ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8) - AdamOptimizerWithDecay(n_epochs, T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ)) -end - -function update!(o::Optimizer{<:AdamOptimizerWithDecay{T}}, C::AdamCache, B::AbstractArray) where T - η = o.method.γ ^ o.step * o.method.η₁ - add!(C.B₁, ((o.method.ρ₁ - o.method.ρ₁^o.step) / (T(1.) - o.method.ρ₁^o.step)) * C.B₁, ((T(1.) - o.method.ρ₁) / (T(1.) - o.method.ρ₁ ^ o.step)) * B) - add!(C.B₂, ((o.method.ρ₂ - o.method.ρ₂^o.step) / (T(1.) - o.method.ρ₂^o.step)) * C.B₂, ((T(1.) - o.method.ρ₂) / (T(1.) - o.method.ρ₂ ^ o.step)) * ⊙²(B)) - mul!(B, -η, /ᵉˡᵉ(C.B₁, scalar_add(racᵉˡᵉ(C.B₂), o.method.δ))) -end \ No newline at end of file diff --git a/src/optimizers/bfgs_cache.jl b/src/optimizers/bfgs_cache.jl deleted file mode 100644 index 1bbcf26b4..000000000 --- a/src/optimizers/bfgs_cache.jl +++ /dev/null @@ -1,56 +0,0 @@ -@doc raw""" - BFGSCache(B) - -Make the cache for the BFGS optimizer based on the array `B`. - -It stores an array for the gradient of the previous time step `B` and the inverse of the Hessian matrix `H`. - -The cache for the inverse of the Hessian is initialized with the idendity. -The cache for the previous gradient information is initialized with the zero vector. - -Note that the cache for `H` is changed iteratively, whereas the cache for `B` is newly assigned at every time step. -""" -struct BFGSCache{T, BT<:AbstractArray{T}, HT<:AbstractMatrix{T}} <: AbstractCache{T} - B::BT - H::HT - function BFGSCache(B::AbstractArray) - zeroB = zero(B) - H_init = initialize_hessian_inverse(zeroB) - new{eltype(B), typeof(zeroB), typeof(H_init)}(zero(B), H_init) - end -end - -@kernel function assign_diagonal_ones_kernel!(B::AbstractMatrix{T}) where T - i = @index(Global) - B[i, i] = one(T) -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::BFGSCache) - show(io, raw"`BFGSCache` that currently stores `B`as ...") - show(io, "text/plain", C.B) - println(io, "") - println(io, "... and `H` as") - show(io, "text/plain", C.H) -end - -# @doc raw""" -# initialize_hessian_inverse(B) -# -# Initialize the inverse of the Hessian for various arrays. -# -# # Implementation -# This requires an implementation of a *vectorization operation* `vec`. This is important for custom arrays. -# """ -function initialize_hessian_inverse(B::AbstractArray{T}) where T - length_of_array = length(vec(B)) - backend = networkbackend(B) - H = KernelAbstractions.zeros(backend, T, length_of_array, length_of_array) - assign_diagonal_ones! = assign_diagonal_ones_kernel!(backend) - assign_diagonal_ones!(H, ndrange=length_of_array) - H -end - -setup_bfgs_cache(ps::NeuralNetworkParameters) = setup_bfgs_cache(params(ps)) -setup_bfgs_cache(ps::NamedTuple) = apply_toNT(setup_bfgs_cache, ps) -setup_bfgs_cache(ps::Tuple) = Tuple([setup_bfgs_cache(x) for x in ps]) -setup_bfgs_cache(B::AbstractArray) = BFGSCache(B) \ No newline at end of file diff --git a/src/optimizers/bfgs_optimizer.jl b/src/optimizers/bfgs_optimizer.jl deleted file mode 100644 index d90f7d740..000000000 --- a/src/optimizers/bfgs_optimizer.jl +++ /dev/null @@ -1,67 +0,0 @@ -@doc raw""" - BFGSOptimizer(η, δ) - -Make an instance of the Broyden-Fletcher-Goldfarb-Shanno (BFGS) optimizer. - -`η` is the *learning rate*. -`δ` is a stabilization parameter. -""" -struct BFGSOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - δ::T - - function BFGSOptimizer(η::T = 1f-2, δ=1f-8) where T - new{T}(η, T(δ)) - end -end - -@doc raw""" - update!(o::Optimizer{<:BFGSOptimizer}, C, B) - -Peform an update with the BFGS optimizer. - -`C` is the cache, `B` contains the gradient information (the output of [`global_rep`](@ref) in general). - -First we compute the *final velocity* with -```julia -vecS = -o.method.η * C.H * vec(B) -``` -and then we update `H` -```julia -C.H .= (𝕀 - ρ * SY) * C.H * (𝕀 - ρ * SY') + ρ * vecS * vecS' -``` -where `SY` is `vecS * Y'` and `𝕀` is the idendity. - -# Implementation - -For stability we use `δ` for computing `ρ`: -```julia -ρ = 1. / (vecS' * Y + o.method.δ) -``` - -This is similar to the [`AdamOptimizer`](@ref) - -# Extended help - -If we have weights on a [`Manifold`](@ref) than the updates are slightly more difficult. -In this case the [`vec`](@ref) operation has to be generalized to the corresponding *global tangent space*. -""" -function update!(o::Optimizer{<:BFGSOptimizer}, C::BFGSCache, B::AbstractArray) - T = eltype(o) - # in the first step we compute the difference between the current and the previous mapped gradients: - Y = vec(B - C.B) - # compute the descent direction - P = -C.H * vec(B) - # compute S - vecS = o.method.η * P - # store gradient - assign!(C.B, copy(B)) - # output final velocity - assign!(vec(B), copy(vecS)) - # compute SY and HY - ρ = one(T) / (vecS' * Y + o.method.δ) - SY = vecS * Y' - 𝕀 = one(SY) - # compute H - C.H .= (𝕀 - ρ * SY) * C.H * (𝕀 - ρ * SY') + ρ * vecS * vecS' -end \ No newline at end of file diff --git a/src/optimizers/default_optimizer.jl b/src/optimizers/default_optimizer.jl deleted file mode 100644 index 622e05fbe..000000000 --- a/src/optimizers/default_optimizer.jl +++ /dev/null @@ -1,3 +0,0 @@ -default_optimizer() = GradientOptimizer() - -default_optimizer(arch::Architecture) = GradientOptimizer() \ No newline at end of file diff --git a/src/optimizers/gradient_optimizer.jl b/src/optimizers/gradient_optimizer.jl deleted file mode 100644 index 46c68365b..000000000 --- a/src/optimizers/gradient_optimizer.jl +++ /dev/null @@ -1,26 +0,0 @@ -@doc raw""" - GradientOptimizer(η) - -Make an instance of a gradient optimizer. - -This is the simplest neural network optimizer. It has no cache and computes the final velocity as: -```math - \mathrm{velocity} \gets - \eta\nabla_\mathrm{weight}L. -``` - -# Implementation - -The operations are done as memory efficiently as possible. -This means the provided ``\nabla_WL`` is mutated via: -```julia -rmul!(∇L, -method.η) -``` -""" -struct GradientOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - GradientOptimizer(η = 1e-2) = new{typeof(η)}(η) -end - -function update!(o::Optimizer{<:GradientOptimizer}, ::GradientCache, B::AbstractVecOrMat) - rmul!(B, -o.method.η) -end \ No newline at end of file diff --git a/src/optimizers/init_optimizer_cache.jl b/src/optimizers/init_optimizer_cache.jl deleted file mode 100644 index c12a6a3f9..000000000 --- a/src/optimizers/init_optimizer_cache.jl +++ /dev/null @@ -1,14 +0,0 @@ -# @doc raw""" -# init_optimizer_cache(method, x) -# -# Initialize the cache corresponding to the weights `x` for a specific method. -# -# # Implementation -# -# Wrapper for the functions `setup_adam_cache`, `setup_momentum_cache`, `setup_gradient_cache`, `setup_bfgs_cache`. -# These appear outside of `optimizer_caches.jl` because the `OptimizerMethods` first have to be defined. -# """ -init_optimizer_cache(::GradientOptimizer, x) = setup_gradient_cache(x) -init_optimizer_cache(::MomentumOptimizer, x) = setup_momentum_cache(x) -init_optimizer_cache(::Union{AdamOptimizer, AdamOptimizerWithDecay}, x) = setup_adam_cache(x) -init_optimizer_cache(::BFGSOptimizer, x) = setup_bfgs_cache(x) \ No newline at end of file diff --git a/src/optimizers/manifold_related/global_sections.jl b/src/optimizers/manifold_related/global_sections.jl deleted file mode 100644 index a3ed00d7d..000000000 --- a/src/optimizers/manifold_related/global_sections.jl +++ /dev/null @@ -1,259 +0,0 @@ -@doc raw""" - GlobalSection(Y) - -Construct a global section for `Y`. - -A global section ``\lambda`` is a mapping from a homogeneous space ``\mathcal{M}`` to the corresponding Lie group ``G`` such that - -```math -\lambda(Y)E = Y, -``` - -Also see [`apply_section`](@ref) and [`global_rep`](@ref). - -# Implementation - -For an implementation of `GlobalSection` for a custom array (especially manifolds), the function [`global_section`](@ref) has to be generalized. -""" -struct GlobalSection{T, AT, λT} - Y::AT - # for now the only lift that is implemented is the Stiefel one - these types will have to be expanded! - λ::λT - - function GlobalSection(Y::AbstractVecOrMat) - λ = global_section(Y) - new{eltype(Y), typeof(Y), typeof(λ)}(Y, λ) - end -end - -function GlobalSection(ps::NamedTuple) - apply_toNT(GlobalSection, ps) -end - -function GlobalSection(ps::NeuralNetworkParameters) - setup_cache(GlobalSection, ps) -end - -@doc raw""" - Matrix(λY::GlobalSection) - -Put `λY` into matrix form. - -This is not recommended if speed is important! - -Use [`apply_section`](@ref) and [`global_rep`](@ref) instead! -""" -function Base.Matrix(λY::GlobalSection) - hcat(Matrix(λY.Y), Matrix(λY.λ)) -end - -@doc raw""" - λY * Y - -Apply the element `λY` onto `Y`. - -Here `λY` is an element of a Lie group and `Y` is an element of a homogeneous space. -""" -Base.:*(λY::GlobalSection, Y::Manifold) = apply_section(λY, Y) - -@doc raw""" - apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT <: StiefelManifold{T}} - -Apply `λY` to `Y₂`. - -Mathematically this is the group action of the element ``\lambda{}Y\in{}G`` on the element ``Y_2`` of the homogeneous space ``\mathcal{M}``. - -Internally it calls [`apply_section!`](@ref). -""" -function apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - Y = StiefelManifold(zero(Y₂.A)) - apply_section!(Y, λY, Y₂) - - Y -end - -@doc raw""" - apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - -Apply `λY` to `Y₂` and store the result in `Y`. - -This is the inplace version of [`apply_section`](@ref). -""" -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:StiefelManifold{T}} - N, n = size(λY.Y) - - @views Y.A .= λY.Y * Y₂.A[1:n, :] + λY.λ * Y₂.A[(n+1):N, :] -end - -function apply_section(λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:GrassmannManifold{T}} - Y = GrassmannManifold(zero(Y₂.A)) - apply_section!(Y, λY, Y₂) - - Y -end - -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AT) where {T, AT<:GrassmannManifold{T}} - N, n = size(λY.Y) - - @views Y.A = λY.Y * Y₂.A[1:n, :] + λY.λ * Y₂.A[(n + 1):N, :] -end - -function apply_section(λY::GlobalSection{T}, Y₂::AbstractVecOrMat{T}) where {T} - Y = copy(Y₂) - apply_section!(Y, λY, Y₂) -end - -function apply_section!(Y::AT, λY::GlobalSection{T, AT}, Y₂::AbstractVecOrMat{T}) where {T, AT<:AbstractVecOrMat{T}} - Y .= Y₂ + λY.Y -end - -function apply_section(λY::NamedTuple, Y₂::NamedTuple) - apply_toNT(apply_section, λY, Y₂) -end - -function apply_section!(Y::NamedTuple, λY::NamedTuple, Y₂::NamedTuple) - apply_toNT(apply_section!, Y, λY, Y₂) -end - -function global_rep(λY::NamedTuple, gx::NamedTuple) - apply_toNT(global_rep, λY, gx) -end - -##auxiliary function -function global_rep(::GlobalSection{T}, gx::AbstractVecOrMat{T}) where {T} - gx -end - -@doc raw""" - global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:StiefelManifold{T}} - -Express `Δ` (an the tangent space of `Y`) as an instance of `StiefelLieAlgHorMatrix`. - -This maps an element from ``T_Y\mathcal{M}`` to an element of ``\mathfrak{g}^\mathrm{hor}``. - -These two spaces are isomorphic where the isomorphism where the isomorphism is established through ``\lambda(Y)\in{}G`` via: - -```math -T_Y\mathcal{M} \to \mathfrak{g}^{\mathrm{hor}}, \Delta \mapsto \lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y). -``` - -Also see [`GeometricMachineLearning.Ω`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: _round -import Random - -Random.seed!(123) - -Y = rand(StiefelManifold, 6, 3) -Δ = rgrad(Y, randn(6, 3)) -λY = GlobalSection(Y) - -_round(global_rep(λY, Δ); digits = 3) - -# output - -6×6 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}: - 0.0 0.679 1.925 0.981 -2.058 0.4 - -0.679 0.0 0.298 -0.424 0.733 -0.919 - -1.925 -0.298 0.0 -1.815 1.409 1.085 - -0.981 0.424 1.815 0.0 0.0 0.0 - 2.058 -0.733 -1.409 0.0 0.0 0.0 - -0.4 0.919 -1.085 0.0 0.0 0.0 -``` - -# Implementation - -The function `global_rep` does in fact not perform the entire map ``\lambda(Y)^{-1}\Omega(Y, \Delta)\lambda(Y)`` but only - -```math -\Delta \mapsto \mathrm{skew}(Y^T\Delta), -``` - -to get the small skew-symmetric matrix ``A\in\mathcal{S}_\mathrm{skew}(n)`` and - -```math -\Delta \mapsto (\lambda(Y)_{[1:N, n:N]}^T \Delta)_{[1:(N-n), 1:n]}, -``` - -to get the arbitrary matrix ``B\in\mathbb{R}^{(N-n)\times{}n}``. -""" -function global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:StiefelManifold{T}} - N, n = size(λY.Y) - StiefelLieAlgHorMatrix( - SkewSymMatrix(λY.Y.A' * Δ), - λY.λ' * Δ, - N, - n - ) -end - -@doc raw""" - global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:GrassmannManifold{T}} - -Express `Δ` (an element of the tangent space of `Y`) as an instance of [`GrassmannLieAlgHorMatrix`](@ref). - -The method `global_rep` for [`GrassmannManifold`](@ref) is similar to that for [`StiefelManifold`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: _round -import Random - -Random.seed!(123) - -Y = rand(GrassmannManifold, 6, 3) -Δ = rgrad(Y, randn(6, 3)) -λY = GlobalSection(Y) - -_round(global_rep(λY, Δ); digits = 3) - -# output - -6×6 GrassmannLieAlgHorMatrix{Float64, Matrix{Float64}}: - 0.0 0.0 0.0 0.981 -2.058 0.4 - 0.0 0.0 0.0 -0.424 0.733 -0.919 - 0.0 0.0 0.0 -1.815 1.409 1.085 - -0.981 0.424 1.815 0.0 0.0 0.0 - 2.058 -0.733 -1.409 0.0 0.0 0.0 - -0.4 0.919 -1.085 0.0 0.0 0.0 -``` -""" -function global_rep(λY::GlobalSection{T, AT}, Δ::AbstractMatrix{T}) where {T, AT<:GrassmannManifold{T}} - N, n = size(λY.Y) - GrassmannLieAlgHorMatrix( - λY.λ' * Δ, - N, - n - ) -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::GlobalSection{T, MT}, B⁽ᵗ⁻¹⁾::AbstractLieAlgHorMatrix{T}, retraction) where {T, MT <: Manifold} - N, n = B⁽ᵗ⁻¹⁾.N, B⁽ᵗ⁻¹⁾.n - expB = retraction(B⁽ᵗ⁻¹⁾) - apply_section!(expB, Λ⁽ᵗ⁻¹⁾, expB) - Λ⁽ᵗ⁻¹⁾.Y.A .= @view expB.A[:, 1:n] - Λ⁽ᵗ⁻¹⁾.λ .= @view expB.A[:, (n+1):N] - - nothing -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::GlobalSection{T, AT}, B⁽ᵗ⁻¹⁾::AT, retraction) where {T, AT <: AbstractVecOrMat{T}} - expB = retraction(B⁽ᵗ⁻¹⁾) - apply_section!(expB, Λ⁽ᵗ⁻¹⁾, expB) - Λ⁽ᵗ⁻¹⁾.Y .= expB - - nothing -end - -function update_section!(Λ⁽ᵗ⁻¹⁾::NamedTuple, B⁽ᵗ⁻¹⁾::NamedTuple, retraction) - update_section_closure!(Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾) = update_section!(Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾, retraction) - apply_toNT(update_section_closure!, Λ⁽ᵗ⁻¹⁾, B⁽ᵗ⁻¹⁾) - - nothing -end \ No newline at end of file diff --git a/src/optimizers/manifold_related/modified_exponential.jl b/src/optimizers/manifold_related/modified_exponential.jl deleted file mode 100644 index 13167391f..000000000 --- a/src/optimizers/manifold_related/modified_exponential.jl +++ /dev/null @@ -1,79 +0,0 @@ -update_algorithm = "while norm(Aⁿ) > ε -mul!(A_temp, Aⁿ, A) -Aⁿ .= A_temp -rmul!(Aⁿ, T(inv(n))) - -𝔄A += Aⁿ -n += 1 -end" - -@doc (raw""" - 𝔄(A) - -Compute ``\mathfrak{A}(A) := \sum_{n=1}^\infty \frac{1}{n!} (A)^{n-1}.`` - -# Implementation - -This uses a Taylor expansion that iteratively adds terms with - -```julia -""" * update_algorithm * raw""" - -``` - -until the norm of `Aⁿ` becomes smaller than machine precision. -The counter `n` in the above algorithm is initialized as `2` -The matrices `Aⁿ` and `𝔄` are initialized as the identity matrix. -""") -function 𝔄(A::AbstractMatrix) - T = eltype(A) - Aⁿ = one(A) - 𝔄A = one(A) - A_temp = zero(A) - n = 2 - ε = eps(T) - while norm(Aⁿ) > ε - mul!(A_temp, Aⁿ, A) - Aⁿ .= A_temp - rmul!(Aⁿ, T(inv(n))) - - 𝔄A += Aⁿ - n += 1 - end - # println("Number of iterations is: ", i) - 𝔄A -end - -@doc raw""" - 𝔄(B̂, B̄) - -Compute ``\mathfrak{A}(B', B'') := \sum_{n=1}^\infty \frac{1}{n!} ((B'')^TB')^{n-1}.`` - -This expression has the property ``\mathbb{I} + B'\mathfrak{A}(B', B'')(B'')^T = \exp(B'(B'')^T).`` - -# Examples - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: 𝔄 -import Random -Random.seed!(123) - -B = rand(StiefelLieAlgHorMatrix, 10, 2) -B̂ = hcat(vcat(.5 * B.A, B.B), vcat(one(B.A), zero(B.B))) -B̄ = hcat(vcat(one(B.A), zero(B.B)), vcat(-.5 * B.A, -B.B)) - -one(B̂ * B̄') + B̂ * 𝔄(B̂, B̄) * B̄' ≈ exp(Matrix(B)) - -# output - -true -``` -""" -function 𝔄(B̂::AbstractMatrix, B̄::AbstractMatrix) - 𝔄(B̄' * B̂) -end - -function 𝔄exp(X::AbstractMatrix{T}, Y::AbstractMatrix{T}) where T - I + X * 𝔄(X, Y) * Y' -end \ No newline at end of file diff --git a/src/optimizers/manifold_related/retraction_types.jl b/src/optimizers/manifold_related/retraction_types.jl deleted file mode 100644 index e38c10bd5..000000000 --- a/src/optimizers/manifold_related/retraction_types.jl +++ /dev/null @@ -1,6 +0,0 @@ -# AbstractRetraction is a type that comprises all retraction methods for manifolds. For every manifold layer one has to specify a retraction method that takes the layer and elements of the (global) tangent space. -abstract type AbstractRetraction end - -struct Cayley <: AbstractRetraction end - -struct Geodesic <: AbstractRetraction end \ No newline at end of file diff --git a/src/optimizers/manifold_related/retractions.jl b/src/optimizers/manifold_related/retractions.jl deleted file mode 100644 index 3dab15aad..000000000 --- a/src/optimizers/manifold_related/retractions.jl +++ /dev/null @@ -1,188 +0,0 @@ -geodesic(A::AbstractVecOrMat) = A -cayley(A::AbstractVecOrMat) = A - -geodesic(B::NamedTuple) = apply_toNT(geodesic, B) - -@doc raw""" - geodesic(Y::Manifold, Δ) - -Take as input an element of a manifold `Y` and a tangent vector in `Δ` in the corresponding tangent space and compute the geodesic (exponential map). - -In different notation: take as input an element ``x`` of ``\mathcal{M}`` and an element of ``T_x\mathcal{M}`` and return ``\mathtt{geodesic}(x, v_x) = \exp(v_x).`` - - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = StiefelManifold([1. 0. 0.;]' |> Matrix) -Δ = [0. .5 0.;]' |> Matrix -Y₂ = geodesic(Y, Δ) - -Y₂' * Y₂ ≈ [1.;] - -# output - -true -``` - -# Implementation - -Internally this `geodesic` method calls [`geodesic(::StiefelLieAlgHorMatrix)`](@ref). -""" -function geodesic(Y::Manifold{T}, Δ::AbstractMatrix{T}) where T - λY = GlobalSection(Y) - B = global_rep(λY, Δ) - E = StiefelProjection(B) - expB = geodesic(B) - λY * typeof(Y)(expB * E) -end - -@doc raw""" - geodesic(B̄::StiefelLieAlgHorMatrix) - -Compute the geodesic of an element in [`StiefelLieAlgHorMatrix`](@ref). - -# Implementation - -Internally this is using: - -```math -\mathbb{I} + B'\mathfrak{A}(B', B'')B'', -``` - -with - -```math -\bar{B} = \begin{bmatrix} - A & -B^T \\ - B & \mathbb{O} -\end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T. -``` - -This is using a computationally efficient version of the matrix exponential ``\mathfrak{A}``. - -See [`GeometricMachineLearning.𝔄`](@ref). -""" -function geodesic(B::StiefelLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - unit = one(B.A) - A_mat = B.A * unit - B̂ = hcat(vcat(T(.5) * A_mat, B.B), E) - B̄ = hcat(vcat(unit, T(.5) * A_mat), vcat(zero(B.B'), -B.B'))' - StiefelManifold(one(B) + B̂ * 𝔄(B̂, B̄) * B̄') -end - -@doc raw""" - geodesic(B̄::GrassmannLieAlgHorMatrix) - -Compute the geodesic of an element in [`GrassmannLieAlgHorMatrix`](@ref). - -This is equivalent to the method of [`geodesic`](@ref) for [StiefelLieAlgHorMatrix](@ref). - -See [`geodesic(::StiefelLieAlgHorMatrix)`](@ref). -""" -function geodesic(B::GrassmannLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - backend = networkbackend(B) - zero_mat = KernelAbstractions.zeros(backend, T, B.n, B.n) - B̂ = hcat(vcat(zero_mat, B.B), E) - B̄ = hcat(vcat(one(zero_mat), zero_mat), vcat(zero(B.B'), -B.B'))' - GrassmannManifold(one(B) + B̂ * 𝔄(B̂, B̄) * B̄') -end - -cayley(B::NamedTuple) = apply_toNT(cayley, B) - -@doc raw""" - cayley(Y::Manifold, Δ) - -Take as input an element of a manifold `Y` and a tangent vector in `Δ` in the corresponding tangent space and compute the Cayley retraction. - -In different notation: take as input an element ``x`` of ``\mathcal{M}`` and an element of ``T_x\mathcal{M}`` and return ``\mathrm{Cayley}(v_x).`` - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = StiefelManifold([1. 0. 0.;]' |> Matrix) -Δ = [0. .5 0.;]' |> Matrix -Y₂ = cayley(Y, Δ) - -Y₂' * Y₂ ≈ [1.;] - -# output - -true -``` - -See the example in [`geodesic(::Manifold{T}, ::AbstractMatrix{T}) where T`]. -""" -function cayley(Y::Manifold{T}, Δ::AbstractMatrix{T}) where T - λY = GlobalSection(Y) - B = global_rep(λY, Δ) - E = StiefelProjection(B) - cayleyB = cayley(B) - λY * typeof(Y)(cayleyB * E) -end - -@doc raw""" - cayley(B̄::StiefelLieAlgHorMatrix) - -Compute the Cayley retraction of `B`. - -# Implementation - -Internally this is using - -```math -\mathrm{Cayley}(\bar{B}) = \mathbb{I} + \frac{1}{2} B' (\mathbb{I}_{2n} - \frac{1}{2} (B'')^T B')^{-1} (B'')^T (\mathbb{I} + \frac{1}{2} B), -``` -with -```math -\bar{B} = \begin{bmatrix} - A & -B^T \\ - B & \mathbb{O} -\end{bmatrix} = \begin{bmatrix} \frac{1}{2}A & \mathbb{I} \\ B & \mathbb{O} \end{bmatrix} \begin{bmatrix} \mathbb{I} & \mathbb{O} \\ \frac{1}{2}A & -B^T \end{bmatrix} =: B'(B'')^T, -``` -i.e. ``\bar{B}`` is expressed as a product of two ``N\times{}2n`` matrices. -""" -function cayley(B::StiefelLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - 𝕀_small = one(B.A) - 𝕆 = zero(𝕀_small) - 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) - 𝕀_big = one(B) - A_mat = B.A * 𝕀_small - B̂ = hcat(vcat(T(.5) * A_mat, B.B), E) - B̄ = hcat(vcat(𝕀_small, T(.5) * A_mat), vcat(zero(B.B'), -B.B'))' - - StiefelManifold((𝕀_big + T(.5) * B̂ * inv(𝕀_small2 - T(.5) * B̄' * B̂) * B̄') * (𝕀_big + T(.5) * B)) -end - -@doc raw""" - cayley(B̄::GrassmannLieAlgHorMatrix) - -Compute the Cayley retraction of `B`. - -This is equivalent to the method of [`cayley`](@ref) for [StiefelLieAlgHorMatrix](@ref). - -See [`cayley(::StiefelLieAlgHorMatrix)`](@ref). -""" -function cayley(B::GrassmannLieAlgHorMatrix) - T = eltype(B) - E = StiefelProjection(B) - backend = networkbackend(B) - 𝕆 = KernelAbstractions.zeros(backend, T, B.n, B.n) - 𝕀_small = one(𝕆) - 𝕀_small2 = hcat(vcat(𝕀_small, 𝕆), vcat(𝕆, 𝕀_small)) - 𝕀_big = one(B) - B̂ = hcat(vcat(𝕆, B.B), E) - B̄ = hcat(vcat(𝕀_small, 𝕆), vcat(zero(B.B'), -B.B'))' - - GrassmannManifold((𝕀_big + T(.5) * B̂ * inv(𝕀_small2 - T(.5) * B̄' * B̂) * B̄') * (𝕀_big + T(.5) * B)) -end \ No newline at end of file diff --git a/src/optimizers/momentum_optimizer.jl b/src/optimizers/momentum_optimizer.jl deleted file mode 100644 index 1404981d3..000000000 --- a/src/optimizers/momentum_optimizer.jl +++ /dev/null @@ -1,32 +0,0 @@ -@doc raw""" - MomentumOptimizer(η, α) - -Make an instance of the momentum optimizer. - -The momentum optimizer is similar to the [`GradientOptimizer`](@ref). -It however has a nontrivial cache that stores past history (see [`MomentumCache`](@ref)). -The cache is updated via: -```math - B^{\mathrm{cache}} \gets \alpha{}B^{\mathrm{cache}} + \nabla_\mathrm{weights}L -``` -and then the final velocity is computed as -```math - \mathrm{velocity} \gets - \eta{}B^{\mathrm{cache}}. -``` - -# Implementation - -To save memory the *velocity* is stored in the input ``\nabla_WL``. -This is similar to the case of the [`GradientOptimizer`](@ref). -""" -struct MomentumOptimizer{T<:Real} <: OptimizerMethod{T} - η::T - α::T - MomentumOptimizer(η = 1e-3, α = 1e-2) = new{typeof(η)}(η, α) -end - -#update for weights -function update!(o::Optimizer{<:MomentumOptimizer}, C::MomentumCache, B::AbstractVecOrMat) - add!(C.B, o.method.α*C.B, B) - mul!(B, -o.method.η, C.B) -end \ No newline at end of file diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl deleted file mode 100644 index 1395ce424..000000000 --- a/src/optimizers/optimizer.jl +++ /dev/null @@ -1,177 +0,0 @@ -@doc raw""" - Optimizer(method, cache, step, retraction) - -Store the `method` (e.g. [`AdamOptimizer`](@ref) with corresponding hyperparameters), the `cache` (e.g. [`AdamCache`](@ref)), the optimization step and the retraction. - -It takes as input an optimization method and the parameters of a network. - -Before one can call `Optimizer` a [`OptimizerMethod`](@ref) that stores all the hyperparameters of the optimizer needs to be specified. - -# Functor - -For an instance `o` of `Optimizer`, we can call the corresponding functor as: - -```julia -o(nn, dl, batch, n_epochs, loss) -``` - -The arguments are: -1. `nn::NeuralNetwork` -2. `dl::`[`DataLoader`](@ref) -3. `batch::`[`Batch`](@ref) -4. `n_epochs::Integer` -5. `loss::NetworkLoss` - -The last argument is optional for many neural network architectures. We have the following defaults: -- A [`TransformerIntegrator`](@ref) uses [`TransformerLoss`](@ref). -- A [`NeuralNetworkIntegrator`](@ref) uses `FeedForwardLoss` (this loss is defined in `AbstractNeuralNetworks`). -- An [`AutoEncoder`](@ref) uses [`AutoEncoderLoss`](@ref). - -In addition there is an optional keyword argument that can be supplied to the functor: -- `show_progress=true`: This specifies whether a progress bar should be shown during training. - -# Implementation - -Internally the functor for `Optimizer` calls [`GlobalSection`](@ref) once at the start and then [`optimize_for_one_epoch!`](@ref) for each epoch. -""" -mutable struct Optimizer{MT<:OptimizerMethod, CT, RT} - method::MT - cache::CT - step::Int - retraction::RT -end - -eltype(::Optimizer{<:OptimizerMethod{T}}) where T = T - -@doc raw""" - Optimizer(method, nn_params) - -Allocate the cache for a specific `method` and `nn_params` for an instance of `Optimizer`. - -Internally this calls [`init_optimizer_cache`](@ref). - -An equivalent constructor is - -```julia -Optimizer(method, nn::NeuralNetwork) -``` - -# Arguments - -The optional keyword argument is the retraction. By default this is [`cayley`](@ref). -""" -function Optimizer(method::OptimizerMethod, nn_params::Union{NeuralNetworkParameters, NamedTuple}; retraction = cayley) - Optimizer(method, init_optimizer_cache(method, nn_params), 0, retraction) -end - -function Optimizer(method::OptimizerMethod, nn::NeuralNetwork; kwargs...) - Optimizer(method, params(nn); kwargs...) -end - -Optimizer(nn::NeuralNetwork, m::OptimizerMethod; kwargs...) = Optimizer(m, nn; kwargs...) - -@doc raw""" - update!(o, cache, B) - -Update the `cache` and output a final velocity that is stored in `B`. - -Note that ``B\in\mathfrak{g}^\mathrm{hor}`` in general. - -In the manifold case the final velocity is the input to a retraction. -""" -function update!(o::Optimizer, ::AbstractCache, ::AbstractArray) - error("No update rule implemented for method", o.method) -end - -####################################################################################### -# optimization step function - -function _optimization_step!(o::Optimizer, λY::NamedTuple, ps::NamedTuple, cache::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) - gx = rgrad(ps, dx) - B = global_rep(λY, gx) - update!(o, cache, B) - update_section!(λY, B, o.retraction) - - nothing -end - -function optimization_step!(o::Optimizer, λY::NamedTuple, ps::NeuralNetworkParameters, dx::Union{NamedTuple, NeuralNetworkParameters}) - @assert keys(o.cache) == keys(λY) == keys(ps) == keys(dx) - o.step += 1 - for key in keys(o.cache) - cache = o.cache[key] - λY_temp = λY[key] - ps_temp = ps[key] - dx_temp = dx[key] - _optimization_step!(o, λY_temp, ps_temp, cache, dx_temp) - end -end - -@doc raw""" - optimization_step!(o, λY, ps, dx) - -Update the weights `ps` based on an [`Optimizer`](@ref), a `cache` and first-order derivatives `dx`. - -`optimization_step!` is calling [`update!`](@ref) internally. -`update!` has to be implemented for every [`OptimizerMethod`](@ref). - -# Arguments - -All arguments into `optimization_step!` are mandatory: -1. `o::`[`Optimizer`](@ref), -2. `λY::NamedTuple`: this named tuple has the same keys as `ps`, but contains [`GlobalSection`](@ref)s, -3. `ps::NamedTuple`: the neural network parameters, -5. `dx::Union{NamedTuple, NeuralNetworkParameters}`: the gradients stores as a NamedTuple. - -All the arguments are given as `NamedTuple`s as the neural network weights are stores in that format. - -```jldoctest -using GeometricMachineLearning -using GeometricMachineLearning: params - -l = StiefelLayer(3, 5) -ps = params(NeuralNetwork(Chain(l), Float32)).L1 -cache = apply_toNT(MomentumCache, ps) -o = Optimizer(MomentumOptimizer(), cache, 0, geodesic) -λY = GlobalSection(ps) -dx = (weight = rand(Float32, 5, 3), ) - -# call the optimizer -optimization_step!(o, λY, ps, dx) - -_test_nt(x) = typeof(x) <: NamedTuple - -_test_nt(λY) & _test_nt(ps) & _test_nt(cache) & _test_nt(dx) - -# output - -true -``` - -# Extended help -The derivatives `dx` here are usually obtained via an AD routine by differentiating a loss function, i.e. `dx` is ``\nabla_xL``. -""" -function optimization_step!(o::Optimizer, λY::NamedTuple, ps::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) - o.step += 1 - - _optimization_step!(o, λY, ps, o.cache, dx) -end - -####################################################################################### -# utils functions (should probably be put somewhere else) - -rgrad(ps::NamedTuple, dx::Union{NamedTuple, NeuralNetworkParameters}) = apply_toNT(rgrad, ps, dx) - -function rgrad(Y::AbstractVecOrMat, dx::AbstractVecOrMat) - @assert size(Y) == size(dx) - dx -end - -# do we need those two? -function update!(m::Optimizer, C::NamedTuple, B::NamedTuple) - apply_toNT(m, C, B, update!) -end - -function apply_toNT(m::Optimizer, ps₁::NamedTuple, ps₂::NamedTuple, fun_name) - apply_toNT((ps₁, ps₂) -> fun_name(m, ps₁, ps₂), ps₁, ps₂) -end diff --git a/src/optimizers/optimizer_caches.jl b/src/optimizers/optimizer_caches.jl deleted file mode 100644 index 9f9dd14ab..000000000 --- a/src/optimizers/optimizer_caches.jl +++ /dev/null @@ -1,120 +0,0 @@ -@doc raw""" - AbstractCache - -`AbstractCache` has subtypes: [`AdamCache`](@ref), [`MomentumCache`](@ref), [`GradientCache`](@ref) and [`BFGSCache`](@ref). - -All of them can be initialized with providing an array (also supporting manifold types). -""" -abstract type AbstractCache{T} end - -############################################################################# -# All the definitions of the caches - -@doc raw""" - AdamCache(Y) - -Store the first and second moment for `Y` (initialized as zeros). - -First and second moments are called `B₁` and `B₂`. - -If the cache is called with an instance of a homogeneous space, e.g. the [`StiefelManifold`](@ref) ``St(n,N)`` it initializes the moments as elements of ``\mathfrak{g}^\mathrm{hor}`` ([`StiefelLieAlgHorMatrix`](@ref)). - -# Examples - -```jldoctest -using GeometricMachineLearning - -Y = rand(StiefelManifold, 5, 3) -AdamCache(Y).B₁ - -# output - -5×5 StiefelLieAlgHorMatrix{Float64, SkewSymMatrix{Float64, Vector{Float64}}, Matrix{Float64}}: - 0.0 -0.0 -0.0 -0.0 -0.0 - 0.0 0.0 -0.0 -0.0 -0.0 - 0.0 0.0 0.0 -0.0 -0.0 - 0.0 0.0 0.0 0.0 0.0 - 0.0 0.0 0.0 0.0 0.0 -``` -""" -struct AdamCache{T, AT <: AbstractArray{T}} <: AbstractCache{T} - B₁::AT - B₂::AT - function AdamCache(Y::AbstractArray) - new{eltype(Y), typeof(zero(Y))}(zero(Y), zero(Y)) - end -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::AdamCache) - println(io, raw"`AdamCache` that currently stores `B₁` as ...") - show(io, "text/plain", C.B₁) - println(io, "") - println(io, raw"and `B₂` as ...") - show(io, "text/plain", C.B₂) -end - -@doc raw""" - MomentumCache(Y) - -Store the moment for `Y` (initialized as zeros). - -The moment is called `B`. - -If the cache is called with an instance of a [`Manifold`](@ref) it initializes the moments as elements of ``\mathfrak{g}^\mathrm{hor}`` ([`AbstractLieAlgHorMatrix`](@ref)). - -See [`AdamCache`](@ref). -""" -struct MomentumCache{T, AT <: AbstractArray{T}} <:AbstractCache{T} - B::AT - function MomentumCache(Y::AbstractArray) - new{eltype(Y), typeof(zero(Y))}(zero(Y)) - end -end - -function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::MomentumCache) - println(io, raw"`MomentumCache` that currently stores `B`as ...") - show(io, "text/plain", C.B) -end - -@doc raw""" - GradientCache(Y) - -Do not store anything. - -The cache for the [`GradientOptimizer`](@ref) does not consider past information. -""" -struct GradientCache{T} <: AbstractCache{T} end -GradientCache(::AbstractArray{T}) where T = GradientCache{T}() - -############################################################################# -# All the setup_cache functions - -setup_adam_cache(ps::NamedTuple) = apply_toNT(setup_adam_cache, ps) -setup_momentum_cache(ps::NamedTuple) = apply_toNT(setup_momentum_cache, ps) -setup_gradient_cache(ps::NamedTuple) = apply_toNT(setup_gradient_cache, ps) - -function setup_cache(_setup_cache_function, ps::NeuralNetworkParameters) - ps_keys = keys(ps) - values = Tuple([_setup_cache_function(ps[key]) for key in ps_keys]) - NamedTuple{ps_keys}(values) -end - -setup_adam_cache(ps::NeuralNetworkParameters) = setup_cache(setup_adam_cache, ps) -setup_momentum_cache(ps::NeuralNetworkParameters) = setup_cache(setup_momentum_cache, ps) -setup_gradient_cache(ps::NeuralNetworkParameters) = setup_cache(setup_gradient_cache, ps) - -setup_adam_cache(B::AbstractArray{<:Number}) = AdamCache(B) -setup_momentum_cache(B::AbstractArray{<:Number}) = MomentumCache(B) -setup_gradient_cache(B::AbstractArray{<:Number}) = GradientCache(B) - -function Base.zero(Y::StiefelManifold{T}) where T - N, n = size(Y) - backend = networkbackend(Y.A) - zeros(backend, StiefelLieAlgHorMatrix{T}, N, n) -end - -function Base.zero(Y::GrassmannManifold{T}) where T - N, n = size(Y) - backend = networkbackend(Y.A) - zeros(backend, GrassmannLieAlgHorMatrix{T}, N, n) -end diff --git a/src/optimizers/optimizer_method.jl b/src/optimizers/optimizer_method.jl deleted file mode 100644 index 290b293c6..000000000 --- a/src/optimizers/optimizer_method.jl +++ /dev/null @@ -1,15 +0,0 @@ -@doc raw""" - OptimizerMethod - -Each `Optimizer` has to be called with an `OptimizerMethod`. This specifies how the neural network weights are updated in each optimization step. -""" -abstract type OptimizerMethod{T} end - -# @doc raw""" -# init_optimizer_cache(method, x) -# -# Initialize the optimizer cache based on input `x` for the given `method`. -# """ -function init_optimizer_cache(::OptimizerMethod, x) end - -eltype(::OptimizerMethod{T}) where T = T \ No newline at end of file diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index e33e081df..d72387f14 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -27,8 +27,10 @@ struct ZygotePullback{NNLT} <: AbstractPullback{NNLT} loss::NNLT end -(_pullback::ZygotePullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps) -(_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) +(_pullback::ZygotePullback)(ps, model, input_nt::QPTOAT)::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_nt), ps) +(_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) """ _processing(returned_pullback) @@ -37,4 +39,4 @@ Strip `returned_pullback` from unnecessary `Zygote`-induces garbage. Also see the docs for [`ZygotePullback`](@ref). """ -_processing = _get_params∘_get_contents \ No newline at end of file +_processing = _get_params ∘ _get_contents diff --git a/src/training/train.jl b/src/training/train.jl index 89ad325ea..bc28b9bfe 100644 --- a/src/training/train.jl +++ b/src/training/train.jl @@ -54,7 +54,7 @@ function train!(nn::AbstractNeuralNetwork, _data::AbstractTrainingData, m::Optim Loss() = typeof(method) <: TrainingMethod ? loss(method, nn, data) : method(nn, data) # creation of optimiser - @timeit to "Creation of Optimizer" opt = Optimizer(m, params(nn)) + @timeit to "Creation of Optimizer" opt = Optimizer(m, nn) # creation of the array to store total loss total_loss = zeros(typeof(Loss()), ntraining) @@ -66,7 +66,7 @@ function train!(nn::AbstractNeuralNetwork, _data::AbstractTrainingData, m::Optim @timeit to "Computing Grad Loss" ∇params = loss_gradient(Loss, index_batch, params(nn)) - @timeit to "Performing Optimization step" optimization_step!(opt, model(nn), params(nn), ∇params) + @timeit to "Performing Optimization step" optimization_step!(opt, GlobalSection(params(nn)), params(nn), ∇params) total_loss[j] = Loss() diff --git a/src/utils.jl b/src/utils.jl index 0fe9b3b59..cba7ae3a3 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,7 +1,7 @@ # Convenient structure struct NothingFunction <: Function end (::NothingFunction)(args...) = nothing -is_NothingFunction(f::Function) = typeof(f)==NothingFunction +is_NothingFunction(f::Function) = typeof(f) == NothingFunction struct UnknownProblem <: AbstractProblem end @@ -9,14 +9,13 @@ const ∞ = Inf # Functions on typple and named tuple -@inline next(i::Int,j::Int) = (i,j+1) -@inline next(i::Int) = (i+1,) +@inline next(i::Int, j::Int) = (i, j + 1) +@inline next(i::Int) = (i + 1,) @inline tuplejoin(x) = x @inline tuplejoin(x, y) = (x..., y...) @inline tuplejoin(x, y, z...) = tuplejoin(tuplejoin(x, y), z...) - rdevelop(x) = x rdevelop(t::Tuple{Any}) = [rdevelop(t[1])...] rdevelop(t::Tuple) = [rdevelop(t[1])..., rdevelop(t[2:end])...] @@ -27,8 +26,7 @@ develop(t::Tuple{Any}) = [develop(t[1])...] develop(t::Tuple) = [develop(t[1])..., develop(t[2:end])...] develop(t::NamedTuple) = vcat([[develop(e)...] for e in t]...) - -_tuplediff(t₁::Tuple,t₂::Tuple) = tuple(setdiff(Set(t₁),Set(t₂))...) +_tuplediff(t₁::Tuple, t₂::Tuple) = tuple(setdiff(Set(t₁), Set(t₂))...) function apply_toNT(fun, ps::NamedTuple...) for p in ps @@ -37,17 +35,23 @@ function apply_toNT(fun, ps::NamedTuple...) NamedTuple{keys(ps[1])}(fun(p...) for p in zip(ps...)) end -# overload norm -_norm(dx::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} = (norm(dx.q) + norm(dx.p)) / √2 # we need this because of a Zygote problem +# overload norm +function _norm(dx::NT) where { + AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} + (norm(dx.q) + norm(dx.p)) / √2 +end # we need this because of a Zygote problem _norm(dx::NamedTuple) = sum(apply_toNT(norm, dx)) / √length(dx) _norm(A::AbstractArray) = norm(A) -# overloaded +/- operation -_diff(dx₁::NT, dx₂::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} = (q = dx₁.q - dx₂.q, p = dx₁.p - dx₂.p) # we need this because of a Zygote problem +# overloaded +/- operation +function _diff(dx₁::NT, + dx₂::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} + (q = dx₁.q - dx₂.q, p = dx₁.p - dx₂.p) +end # we need this because of a Zygote problem _diff(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_diff, dx₁, dx₂) -_diff(A::AbstractArray, B::AbstractArray) = A - B +_diff(A::AbstractArray, B::AbstractArray) = A - B _add(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_add, dx₁, dx₂) -_add(A::AbstractArray, B::AbstractArray) = A + B +_add(A::AbstractArray, B::AbstractArray) = A + B function add!(C::AbstractVecOrMat, A::AbstractVecOrMat, B::AbstractVecOrMat) @assert size(A) == size(B) == size(C) @@ -61,19 +65,18 @@ end # Type pyracy!! function Base.:+(a::Float64, b::Tuple{Float64}) x, = b - return a+x + return a + x end # Type pyracy!! function Base.:+(a::Vector{Float64}, b::Tuple{Float64}) x, = b y, = a - return y+x + return y + x end - -# Kernel that is needed for functions relating to `SymmetricMatrix` and `SkewSymMatrix` -@kernel function write_ones_kernel!(unit_matrix::AbstractMatrix{T}) where T +# Kernel that is needed for functions relating to `SymmetricMatrix` and `SkewSymMatrix` +@kernel function write_ones_kernel!(unit_matrix::AbstractMatrix{T}) where {T} i = @index(Global) unit_matrix[i, i] = one(T) end @@ -92,18 +95,17 @@ end # utils functions on string function type_without_brace(var) type_str = string(typeof(var)) - replace(type_str, r"\{.*\}"=>"") + replace(type_str, r"\{.*\}" => "") end -function center_align_text(text,width) +function center_align_text(text, width) padding = max(0, width - length(text)) - left_padding = repeat(" ",padding ÷2) + left_padding = repeat(" ", padding ÷ 2) right_padding = repeat(" ", padding - length(left_padding)) aligned_text = left_padding * text * right_padding return aligned_text end - # The following are fallback functions - maybe you want to put them into a separate file function global_section(::AbstractVecOrMat) nothing @@ -163,10 +165,216 @@ const QPTOAT = Union{QPT, AbstractArray} This could be data in ``(q, p)\in\mathbb{R}^{2d}`` form or come from an arbitrary vector space. """ -const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where T +const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where {T} Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p) _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) -_eltype(ps::Tuple) = _eltype(ps[1]) \ No newline at end of file +_eltype(ps::Tuple) = _eltype(ps[1]) +_eltype(ps::NeuralNetworkParameters) = _eltype(params(ps)[1]) + +# Extend GlobalSection so it works with NeuralNetworkParameters (wraps a NamedTuple). +GeometricOptimizers.GlobalSection(ps::NeuralNetworkParameters) = + GeometricOptimizers.GlobalSection(params(ps)) + +# Backward-compat alias +const AbstractCache{T} = GeometricOptimizers.OptimizerCache{T} + +# Gradient wrapper: stores a pre-computed Euclidean gradient and applies rgrad on manifolds. +mutable struct _GMLGradient{T, VT} <: GeometricOptimizers.Gradient{T} + dp::VT +end +(g::_GMLGradient{T})(x::GeometricOptimizers.ArrayNamedTuple{T}) where {T} = + GeometricOptimizers.apply_toNT(rgrad, x, g.dp) +(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp + +# State for Euclidean (non-manifold) parameters. +mutable struct GMLEuclideanState{T, AT<:AbstractArray{T}} + iterations::Int + m₁::AT + m₂::AT +end +GMLEuclideanState(x::AbstractArray{T}) where T = + GMLEuclideanState{T, typeof(x)}(0, zero(x), zero(x)) + +# Adam with exponential learning-rate decay. +struct AdamOptimizerWithDecay{T<:Real} <: GeometricOptimizers.OptimizerMethod + η₁::T; η₂::T; ρ₁::T; ρ₂::T; δ::T; γ::T; n_epochs::Int + function AdamOptimizerWithDecay(n_epochs::Int, η₁=1f-2, η₂=1f-6, + ρ₁=9f-1, ρ₂=9.9f-1, δ=1f-8; T=typeof(η₁)) + γ = exp(log(η₂/η₁) / n_epochs) + new{T}(T(η₁), T(η₂), T(ρ₁), T(ρ₂), T(δ), T(γ), n_epochs) + end +end + +_is_go_native_method(::GeometricOptimizers.GradientMethod) = true +_is_go_native_method(::GeometricOptimizers.MomentumMethod) = true +_is_go_native_method(::GeometricOptimizers.Adam) = true +_is_go_native_method(::GeometricOptimizers.OptimizerMethod) = false + +_adapt_method_to_T(method::GeometricOptimizers.Adam, ::Type{T}) where T = + GeometricOptimizers.Adam(T(method.η), T(method.β₁), T(method.β₂), T(method.δ)) +_adapt_method_to_T(method::GeometricOptimizers.MomentumMethod, ::Type{T}) where T = + GeometricOptimizers.MomentumMethod(T(method.α)) +_adapt_method_to_T(method, ::Type) = method + +_use_go_cache(method, x) = + _is_go_native_method(method) && x isa GeometricOptimizers.OptimizerSolution + +function _make_optimizer_cache(method, x) + if _use_go_cache(method, x) + GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, _eltype(x)), x) + elseif x isa NamedTuple || x isa NeuralNetworkParameters + NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x))) + else + GMLEuclideanState(x) + end +end + +function _make_optimizer_state(method, x) + if _use_go_cache(method, x) + GeometricOptimizers.OptimizerState(method, x) + elseif x isa NamedTuple || x isa NeuralNetworkParameters + NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x))) + else + GMLEuclideanState(x) + end +end + +mutable struct Optimizer{MT <: GeometricOptimizers.OptimizerMethod, CT, ST, RT} + method::MT + cache::CT + state::ST + retraction::RT + step_size::Float64 + iterations::Int +end + +_default_step_size(method::GeometricOptimizers.Adam) = Float64(method.η) +_default_step_size(method::AdamOptimizerWithDecay) = Float64(method.η₁) +_default_step_size(::GeometricOptimizers.OptimizerMethod) = 1e-2 + +_current_step_size(opt::Optimizer, ::Int) = opt.step_size +_current_step_size(opt::Optimizer{<:AdamOptimizerWithDecay}, t::Int) = + Float64(opt.method.η₁ * opt.method.γ^t) + +function Optimizer(method::GeometricOptimizers.OptimizerMethod, nn::NeuralNetwork; + retraction = GeometricOptimizers.cayley, + step_size::Real = _default_step_size(method)) + ps = params(nn) + Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), + retraction, Float64(step_size), 0) +end + +function Optimizer(method::GeometricOptimizers.OptimizerMethod, + ps::Union{NamedTuple, NeuralNetworkParameters}; + retraction = GeometricOptimizers.cayley, + step_size::Real = _default_step_size(method)) + Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), + retraction, Float64(step_size), 0) +end + +# Euclidean update rules +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState, ::GeometricOptimizers.GradientMethod, step_size) where T + x .-= T(step_size) .* dx +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::GeometricOptimizers.MomentumMethod, step_size) where T + x .-= T(step_size) .* (dx .+ state.m₁) + state.m₁ .+= T(method.α) .* dx +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::GeometricOptimizers.Adam, step_size) where T + t = state.iterations; _t = t + 1 + β₁, β₂, δ = T(method.β₁), T(method.β₂), T(method.δ) + fac₁₁ = β₁/(1-β₁^_t); fac₁₂ = (1-β₁)/(1-β₁^_t) + fac₂₁ = β₂/(1-β₂^_t); fac₂₂ = (1-β₂)/(1-β₂^_t) + state.m₁ .= fac₁₁ .* state.m₁ .+ fac₁₂ .* dx + state.m₂ .= fac₂₁ .* state.m₂ .+ fac₂₂ .* dx .^ 2 + x .-= T(step_size) .* state.m₁ ./ (sqrt.(state.m₂) .+ δ) +end +function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, + state::GMLEuclideanState{T}, method::AdamOptimizerWithDecay, step_size) where T + t = state.iterations; _t = t + 1 + ρ₁, ρ₂, δ = T(method.ρ₁), T(method.ρ₂), T(method.δ) + fac₁₁ = ρ₁/(1-ρ₁^_t); fac₁₂ = (1-ρ₁)/(1-ρ₁^_t) + fac₂₁ = ρ₂/(1-ρ₂^_t); fac₂₂ = (1-ρ₂)/(1-ρ₂^_t) + state.m₁ .= fac₁₁ .* state.m₁ .+ fac₁₂ .* dx + state.m₂ .= fac₂₁ .* state.m₂ .+ fac₂₂ .* dx .^ 2 + x .-= T(step_size) .* state.m₁ ./ (sqrt.(state.m₂) .+ δ) +end + +# GO-managed leaf step (manifolds, vectors, ArrayNamedTuples) +function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, + state::GeometricOptimizers.OptimizerState, + dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) + T = _eltype(ps_leaf) + local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) + adapted = _adapt_method_to_T(method, T) + if adapted isa GeometricOptimizers.Adam + GeometricOptimizers.update!(cache, state, local_grad, adapted, ps_leaf) + else + GeometricOptimizers.update!(cache, state, local_grad, + GeometricOptimizers.NoHessian{T}(), ps_leaf) + end + GeometricOptimizers._rmul!(GeometricOptimizers.direction(cache), step_size) + GeometricOptimizers.update_section!(GeometricOptimizers.section(cache), + GeometricOptimizers.section(state), + GeometricOptimizers.direction(cache), + retraction) + GeometricOptimizers._copyto!(GeometricOptimizers.solution(cache), + GeometricOptimizers.section(cache)) + GeometricOptimizers._copyto!(ps_leaf, GeometricOptimizers.solution(cache)) + GeometricOptimizers._copyto!(λY_leaf, GeometricOptimizers.section(cache)) + GeometricOptimizers.update_section!(GeometricOptimizers.section(state), + GeometricOptimizers.section(state), + GeometricOptimizers.direction(cache), + retraction) + if state isa GeometricOptimizers.AdamState + GeometricOptimizers._copyto!(GeometricOptimizers.first_moment(state), + GeometricOptimizers.first_moment(cache)) + GeometricOptimizers._copyto!(GeometricOptimizers.second_moment(state), + GeometricOptimizers.second_moment(cache)) + elseif state isa GeometricOptimizers.MomentumState + GeometricOptimizers._add!(GeometricOptimizers.momentum(state), + GeometricOptimizers._mul(adapted.α, + GeometricOptimizers.gradient_array(cache))) + end + state.iterations += 1 + nothing +end + +# Euclidean leaf step (plain AbstractArray params) +function _leaf_optim_step!(cache::GMLEuclideanState, state::GMLEuclideanState, + dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) + _euclidean_update!(ps_leaf, dp_leaf, state, method, step_size) + state.iterations += 1 + nothing +end + +# Recursive dispatcher over the parameter tree +function _tree_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) + if caches isa NamedTuple + for k in keys(caches) + dp_k = dp[k] + dp_k === nothing && continue + λY_k = λY isa NamedTuple ? λY[k] : λY + _tree_optim_step!(caches[k], states[k], dp_k, ps[k], λY_k, + method, retraction, step_size) + end + else + _leaf_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) + end + nothing +end + +function optimization_step!(opt::Optimizer, λY, ps, dp) + step = _current_step_size(opt, opt.iterations) + _tree_optim_step!(opt.cache, opt.state, dp, ps, λY, opt.method, opt.retraction, step) + opt.iterations += 1 + nothing +end + +check(::Optimizer) = nothing diff --git a/test/data_loader/data_loader_for_input_and_output.jl b/test/data_loader/data_loader_for_input_and_output.jl index 3e81aa4b4..0ee3c7f29 100644 --- a/test/data_loader/data_loader_for_input_and_output.jl +++ b/test/data_loader/data_loader_for_input_and_output.jl @@ -16,7 +16,7 @@ nepochs = 2000 model = Chain(Dense(1, nwidth), Dense(nwidth, 1)) nn = NeuralNetwork(model, Float32) dl = DataLoader(xsamples, ysamples) -o = Optimizer(AdamOptimizer(), nn) +o = Optimizer(Adam(), nn) batch = Batch(nbatch, 1, 1) loss = FeedForwardLoss() diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index 22a5c5c46..283b63a87 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -25,13 +25,14 @@ function test_optimization_with_adam(;T=Float32, dim₁=6, dim₂=6, n_images=10 # input dim is dim₁ / patch_length * dim₂ / pach_length; the transformer is called with dim₁ / patch_length and two layers model = Chain(Transformer(dl.input_dim, patch_length, 2; Stiefel=true), ClassificationLayer(dl.input_dim, 10, σ)) - ps = NeuralNetwork(model, CPU(), Float32).params + nn_obj = NeuralNetwork(model, CPU(), Float32) + ps = nn_obj.params loss = FeedForwardLoss() loss₁ = loss(model, ps, dl.input, dl.output) - opt = Optimizer(AdamOptimizer(), ps) + opt = Optimizer(Adam(), nn_obj) λY = GlobalSection(ps) loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) diff --git a/test/network_losses/losses_and_optimization.jl b/test/network_losses/losses_and_optimization.jl index 82f14e5f7..9ace7aca4 100644 --- a/test/network_losses/losses_and_optimization.jl +++ b/test/network_losses/losses_and_optimization.jl @@ -17,7 +17,7 @@ function train_network(; n_epochs=10) nn = setup_network(dl) loss = FeedForwardLoss() - o = Optimizer(AdamOptimizer(), nn) + o = Optimizer(Adam(), nn) batch = Batch(5, 1) loss_array = o(nn, dl, batch, n_epochs, loss) T = eltype(dl) diff --git a/test/optimizers/bfgs_optimizer.jl b/test/optimizers/bfgs_optimizer.jl index 4f34bd550..e4bdb05fb 100644 --- a/test/optimizers/bfgs_optimizer.jl +++ b/test/optimizers/bfgs_optimizer.jl @@ -1,5 +1,4 @@ using GeometricMachineLearning -using GeometricMachineLearning: update_section! using Zygote using Test using LinearAlgebra: norm @@ -8,87 +7,44 @@ import Random Random.seed!(123) @doc raw""" - bfgs_optimizer(N) - -Test if BFGS optimizer perfroms better than gradient optimizer. - -The test is performed on a simple loss function -```math - \mathrm{loss}(A) = norm(A - B) ^ 3, -``` -where ``B`` is fixed. +Test that gradient descent reduces loss on a Euclidean (plain matrix) problem. +BFGSOptimizer is no longer available; this file retains the structure for future extension. """ -function bfgs_optimizer(N; n_steps = 10, η = 1e-4) +function gradient_optimizer_euclidean(N; n_steps = 20, step_size = 1e-3) B = inv(rand(N, N)) - loss(A) = norm(A - B) ^ (2) - A = randn(N, N) - loss1 = loss(A) - method₁ = GradientOptimizer(η) - o₁ = Optimizer(method₁, (A = A,)) - for _ in 1:n_steps - ∇L = Zygote.gradient(loss, A)[1] - update!(o₁, o₁.cache.A, ∇L) - A .+= ∇L - end - loss2 = loss(A) + loss(ps) = norm(ps.A - B) ^ 2 A = randn(N, N) - method₂ = BFGSOptimizer(η) - o₂ = Optimizer(method₂, (A = A,)) + ps = (A = A,) + loss1 = loss(ps) + o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps - ∇L = Zygote.gradient(loss, A)[1] - update!(o₂, o₂.cache.A, ∇L) - A .+= ∇L + ∇L = Zygote.gradient(loss, ps)[1] + λY = GlobalSection(ps) + optimization_step!(o, λY, ps, ∇L) end - loss3 = loss(A) - @test loss1 > loss2 > loss3 - println(loss2) - println(loss3) - + loss2 = loss(ps) + @test loss1 > loss2 end -bfgs_optimizer(10) - @doc raw""" - bfgs_optimizer(N, n) - -Test if BFGS optimizer perfroms better than gradient optimizer. - -The test is performed on a simple loss function -```math - \mathrm{loss}(A) = norm(AA^T - B) ^ 3, -``` -where ``B = Y_BY_B^T`` for some ``Y\in{}St(n, N)`` is fixed. -``A`` in the equation above is optimized on the Stiefel manifold. +Test that gradient descent reduces loss on the Stiefel manifold. """ -function bfgs_stiefel_optimizer(N, n; n_steps = 10, η = 1e-4) +function gradient_optimizer_stiefel(N, n; n_steps = 20, step_size = 1e-3) YB = rand(StiefelManifold, N, n) - B = YB * YB' - loss(A) = norm(A * A' - B) ^ 2 - Y = rand(StiefelManifold, N, n) - λY = GlobalSection(Y) - loss1 = loss(Y) - method₁ = GradientOptimizer(η) - o₁ = Optimizer(method₁, (A = Y,)) - for _ in 1:n_steps - ∇L = Zygote.gradient(loss, Y)[1] - gradL = global_rep(λY, ∇L) - update!(o₁, o₁.cache.A, gradL) - update_section!(λY, gradL, cayley) - end - loss2 = loss(Y) - Y = rand(StiefelManifold, N, n) - method₂ = BFGSOptimizer(η) - o₂ = Optimizer(method₂, (A = Y,)) + B = YB * YB' + loss(ps) = norm(ps.Y * ps.Y' - B) ^ 2 + Y = rand(StiefelManifold, N, n) + ps = (Y = Y,) + loss1 = loss(ps) + o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps - ∇L = Zygote.gradient(loss, Y)[1] - gradL = global_rep(λY, ∇L) - update!(o₂, o₂.cache.A, gradL) - update_section!(λY, gradL, cayley) + ∇L = Zygote.gradient(loss, ps)[1] + λY = GlobalSection(ps) + optimization_step!(o, λY, ps, ∇L) end - loss3 = loss(Y) - @test loss1 > loss2 > loss3 - println(loss2) - println(loss3) + loss2 = loss(ps) + @test loss1 > loss2 end -bfgs_stiefel_optimizer(10, 5) \ No newline at end of file +gradient_optimizer_euclidean(10) +gradient_optimizer_stiefel(10, 5) diff --git a/test/optimizers/optimizer_convergence_tests/psd_optim.jl b/test/optimizers/optimizer_convergence_tests/psd_optim.jl index 6135a83c2..c926a2154 100644 --- a/test/optimizers/optimizer_convergence_tests/psd_optim.jl +++ b/test/optimizers/optimizer_convergence_tests/psd_optim.jl @@ -38,9 +38,9 @@ function svd_test(A, n, train_steps=1000, tol=1e-1; retraction=cayley) model = Chain(PSDLayer(2*N, 2*n), PSDLayer(2*n, 2*N)) ps = NeuralNetwork(model, CPU(), Float64).params - o₁ = Optimizer(GradientOptimizer(0.01), ps; retraction = retraction) - o₂ = Optimizer(MomentumOptimizer(0.01), ps; retraction = retraction) - o₃ = Optimizer(AdamOptimizer(0.01), ps; retraction = retraction) + o₁ = Optimizer(GradientMethod(), ps; retraction = retraction, step_size = 0.01) + o₂ = Optimizer(MomentumMethod(), ps; retraction = retraction, step_size = 0.01) + o₃ = Optimizer(Adam(0.01), ps; retraction = retraction) U₁, Ũ₁, err₁ = train_network!(o₁, model, deepcopy(ps), A, train_steps, tol) U₂, Ũ₂, err₂ = train_network!(o₂, model, deepcopy(ps), A, train_steps, tol) diff --git a/test/optimizers/optimizer_convergence_tests/svd_optim.jl b/test/optimizers/optimizer_convergence_tests/svd_optim.jl index e260c40ff..378bd002b 100644 --- a/test/optimizers/optimizer_convergence_tests/svd_optim.jl +++ b/test/optimizers/optimizer_convergence_tests/svd_optim.jl @@ -34,9 +34,9 @@ function svd_test(A, n, train_steps=1000, tol=1e-1; retraction=cayley) model = Chain(StiefelLayer(N, n), StiefelLayer(n, N)) ps = NeuralNetwork(model, CPU(), Float64).params - o₁ = Optimizer(GradientOptimizer(0.01), ps; retraction = retraction) - o₂ = Optimizer(MomentumOptimizer(0.01), ps; retraction = retraction) - o₃ = Optimizer(AdamOptimizer(0.01), ps; retraction = retraction) + o₁ = Optimizer(GradientMethod(), ps; retraction = retraction, step_size = 0.01) + o₂ = Optimizer(MomentumMethod(), ps; retraction = retraction, step_size = 0.01) + o₃ = Optimizer(Adam(0.01), ps; retraction = retraction) U₁, Ũ₁, err₁ = train_network!(o₁, model, deepcopy(ps), A, train_steps, tol) U₂, Ũ₂, err₂ = train_network!(o₂, model, deepcopy(ps), A, train_steps, tol) diff --git a/test/runtests.jl b/test/runtests.jl index 665d5b11f..2bba056b1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,85 +1,206 @@ using SafeTestsets, Test, GeometricMachineLearning using Documenter: doctest -# @testset "Doc tests " begin doctest(GeometricMachineLearning; manual = false) end +@testset "Doc tests " begin + doctest(GeometricMachineLearning; manual=false) +end # reduced order modeling tests -@safetestset "PSD tests " begin include("psd_architecture_tests.jl") end -@safetestset "SymplecticAutoencoder tests " begin include("symplectic_autoencoder_tests.jl") end -@safetestset "Check if autoencoder error is lower than PSD error " begin include("sae_error_lower_than_psd_error.jl") end -@safetestset "Check reduced model " begin include("reduced_system.jl") end - -@safetestset "Check parameterlength " begin include("parameterlength/check_parameterlengths.jl") end - -@safetestset "Arrays #1 " begin include("arrays/array_tests.jl") end -@safetestset "Map to skew " begin include("arrays/map_to_skew.jl") end -@safetestset "Sampling of arrays " begin include("arrays/random_generation_of_custom_arrays.jl") end -@safetestset "Addition tests for custom arrays " begin include("arrays/addition_tests_for_custom_arrays.jl") end -@safetestset "Scalar multiplication tests for custom arrays " begin include("arrays/scalar_multiplication_for_custom_arrays.jl") end -@safetestset "Matrix multiplication tests for custom arrays " begin include("arrays/matrix_multiplication_for_custom_arrays.jl") end -@safetestset "Test constructors for custom arrays " begin include("arrays/constructor_tests_for_custom_arrays.jl") end -@safetestset "Symplectic Potential (array tests) " begin include("arrays/poisson_tensor.jl") end -@safetestset "Test StiefelLieAlgHorMatrix constructors and lifts " begin include("arrays/test_stiefel_lie_alg_hor_constructors.jl") end -@safetestset "Test GrassmannLieAlgHorMatrix constructors and lifts " begin include("arrays/test_grassmann_lie_alg_hor_constructors.jl") end - -@safetestset "Test triangular matrices " begin include("arrays/triangular.jl") end -@safetestset "Manifolds (Stiefel): " begin include("manifolds/stiefel_manifold.jl") end -@safetestset "Manifolds (Grassmann): " begin include("manifolds/grassmann_manifold.jl") end -@safetestset "Gradient Layer " begin include("layers/gradient_layer_tests.jl") end -@safetestset "Test symplecticity of upscaling layer " begin include("layers/sympnet_layers_test.jl") end -@safetestset "Hamiltonian Neural Network " begin include("hamiltonian_neural_network_tests.jl") end -@safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end - -@safetestset "Custom tensor matrix multiplication " begin include("kernels/tensor_mat_mul.jl") end -@safetestset "Custom inverse for 2x2, 3x3, 4x4, 5x5 matrices " begin include("kernels/tensor_inverse.jl") end -@safetestset "Custom AD rules for kernels " begin include("custom_ad_rules/kernel_pullbacks.jl") end -@safetestset "ResNet " begin include("layers/resnet_tests.jl") end - +@safetestset "PSD tests " begin + include("psd_architecture_tests.jl") +end +@safetestset "SymplecticAutoencoder tests " begin + include("symplectic_autoencoder_tests.jl") +end +@safetestset "Check if autoencoder error is lower than PSD error " begin + include("sae_error_lower_than_psd_error.jl") +end +@safetestset "Check reduced model " begin + include("reduced_system.jl") +end +@safetestset "Check parameterlength " begin + include("parameterlength/check_parameterlengths.jl") +end +@safetestset "Arrays #1 " begin + include("arrays/array_tests.jl") +end +@safetestset "Map to skew " begin + include("arrays/map_to_skew.jl") +end +@safetestset "Sampling of arrays " begin + include("arrays/random_generation_of_custom_arrays.jl") +end +@safetestset "Addition tests for custom arrays " begin + include("arrays/addition_tests_for_custom_arrays.jl") +end +@safetestset "Scalar multiplication tests for custom arrays " begin + include("arrays/scalar_multiplication_for_custom_arrays.jl") +end +@safetestset "Matrix multiplication tests for custom arrays " begin + include("arrays/matrix_multiplication_for_custom_arrays.jl") +end +@safetestset "Test constructors for custom arrays " begin + include("arrays/constructor_tests_for_custom_arrays.jl") +end +@safetestset "Symplectic Potential (array tests) " begin + include("arrays/poisson_tensor.jl") +end +@safetestset "Test StiefelLieAlgHorMatrix constructors and lifts " begin + include("arrays/test_stiefel_lie_alg_hor_constructors.jl") +end +@safetestset "Test GrassmannLieAlgHorMatrix constructors and lifts " begin + include("arrays/test_grassmann_lie_alg_hor_constructors.jl") +end +@safetestset "Test triangular matrices " begin + include("arrays/triangular.jl") +end +@safetestset "Manifolds (Stiefel): " begin + include("manifolds/stiefel_manifold.jl") +end +@safetestset "Manifolds (Grassmann): " begin + include("manifolds/grassmann_manifold.jl") +end +@safetestset "Gradient Layer " begin + include("layers/gradient_layer_tests.jl") +end +@safetestset "Test symplecticity of upscaling layer " begin + include("layers/sympnet_layers_test.jl") +end +@safetestset "Hamiltonian Neural Network " begin + include("hamiltonian_neural_network_tests.jl") +end +@safetestset "Manifold Neural Network Layers " begin + include("layers/manifold_layers.jl") +end +@safetestset "Custom tensor matrix multiplication " begin + include("kernels/tensor_mat_mul.jl") +end +@safetestset "Custom inverse for 2x2, 3x3, 4x4, 5x5 matrices " begin + include("kernels/tensor_inverse.jl") +end +@safetestset "Custom AD rules for kernels " begin + include("custom_ad_rules/kernel_pullbacks.jl") +end +@safetestset "ResNet " begin + include("layers/resnet_tests.jl") +end # transformer-related tests -@safetestset "Test setup of MultiHeadAttention layer Stiefel weights " begin include("transformer_related/multi_head_attention_stiefel_setup.jl") end -@safetestset "Test geodesic and Cayley retr for the MultiHeadAttention layer w/ St weights " begin include("transformer_related/multi_head_attention_stiefel_retraction.jl") end -@safetestset "Test the correct setup of the various optimizer caches for MultiHeadAttention " begin include("transformer_related/multi_head_attention_stiefel_optim_cache.jl") end -@safetestset "Check if the transformer can be applied to a tensor. " begin include("transformer_related/transformer_application.jl") end -@safetestset "Check if the gradient/pullback of MultiHeadAttention changes type in St case " begin include("transformer_related/transformer_gradient.jl") end -@safetestset "Check if the optimization_step! changes the parameters of the transformer " begin include("transformer_related/transformer_optimizer.jl") end - -@safetestset "Attention layer #1 " begin include("attention_layer/attention_setup.jl") end -@safetestset "Classification layer " begin include("layers/classification.jl") end -@safetestset "Optimizer #1 " begin include("optimizers/utils/global_sections.jl") end -@safetestset "Optimizer #2 " begin include("optimizers/utils/optimization_step.jl") end -@safetestset "Optimizer #3 " begin include("optimizers/utils/modified_exponential.jl") end -@safetestset "Optimizer #4 " begin include("optimizers/optimizer_convergence_tests/svd_optim.jl") end -@safetestset "Optimizer #5 " begin include("optimizers/optimizer_convergence_tests/psd_optim.jl") end -@safetestset "Check if Adam with decay converges " begin include("optimizers/optimizer_convergence_tests/adam_with_learning_rate_decay.jl") end -@safetestset "BFGS Optimizer tests " begin include("optimizers/bfgs_optimizer.jl") end -@safetestset "Data " begin include("data/test_data.jl") end -@safetestset "Batch " begin include("data/test_batch.jl") end +@safetestset "Test setup of MultiHeadAttention layer Stiefel weights " begin + include("transformer_related/multi_head_attention_stiefel_setup.jl") +end +@safetestset "Test geodesic and Cayley retr for the MultiHeadAttention layer w/ St weights " begin + include("transformer_related/multi_head_attention_stiefel_retraction.jl") +end +@safetestset "Test the correct setup of the various optimizer caches for MultiHeadAttention " begin + include("transformer_related/multi_head_attention_stiefel_optim_cache.jl") +end +@safetestset "Check if the transformer can be applied to a tensor. " begin + include("transformer_related/transformer_application.jl") +end +@safetestset "Check if the gradient/pullback of MultiHeadAttention changes type in St case " begin + include("transformer_related/transformer_gradient.jl") +end +@safetestset "Check if the optimization_step! changes the parameters of the transformer " begin + include("transformer_related/transformer_optimizer.jl") +end + +@safetestset "Attention layer #1 " begin + include("attention_layer/attention_setup.jl") +end +@safetestset "Classification layer " begin + include("layers/classification.jl") +end +@safetestset "Optimizer #1 " begin + include("optimizers/utils/global_sections.jl") +end +@safetestset "Optimizer #2 " begin + include("optimizers/utils/optimization_step.jl") +end +@safetestset "Optimizer #3 " begin + include("optimizers/utils/modified_exponential.jl") +end +@safetestset "Optimizer #4 " begin + include("optimizers/optimizer_convergence_tests/svd_optim.jl") +end +@safetestset "Optimizer #5 " begin + include("optimizers/optimizer_convergence_tests/psd_optim.jl") +end +@safetestset "Check if Adam with decay converges " begin + include("optimizers/optimizer_convergence_tests/adam_with_learning_rate_decay.jl") +end +@safetestset "BFGS Optimizer tests " begin + include("optimizers/bfgs_optimizer.jl") +end +@safetestset "Data " begin + include("data/test_data.jl") +end +@safetestset "Batch " begin + include("data/test_batch.jl") +end # @safetestset "Method " begin include("train!/test_method.jl") end -@safetestset "Matching " begin include("data/test_matching.jl") end - -@safetestset "Test data loader for q and p data " begin include("data_loader/batch_data_loader_qp_test.jl") end -@safetestset "Test mnist_utils. " begin include("data_loader/mnist_utils.jl") end -@safetestset "Test the data loader in combination with optimization_step! " begin include("data_loader/data_loader_optimization_step.jl") end -@safetestset "Optimizer functor with data loader for Adam " begin include("data_loader/optimizer_functor_with_adam.jl") end -@safetestset "Test data loader for a tensor (q and p data) " begin include("data_loader/draw_batch_for_tensor_test.jl") end - -@safetestset "Test NetworkLoss + Optimizer " begin include("network_losses/losses_and_optimization.jl") end - -@safetestset "Test parallel inverses " begin include("kernels/tensor_inverse.jl") end -@safetestset "Test parallel Cayley " begin include("kernels/tensor_cayley.jl") end - -@safetestset "Test volume-preserving feedforward neural network " begin include("layers/volume_preserving_feedforward.jl") end - -@safetestset "SympNet integrator " begin include("sympnet_integrator.jl") end -@safetestset "Regular transformer integrator " begin include("standard_transformer_integrator.jl") end - -@safetestset "Batch functor(s) " begin include("batch/batch_functor.jl") end - -@safetestset "Volume-Preserving Transformer (skew-symmetric tests) " begin include("volume_preserving_attention/test_skew_map.jl") end -@safetestset "Volume-Preserving Transformer (cayley-transform tests) " begin include("volume_preserving_attention/test_cayley_transforms.jl") end - -@safetestset "Linear Symplectic Attention " begin include("linear_symplectic_attention.jl") end -@safetestset "Linear Symplectic Transformer " begin include("linear_symplectic_transformer.jl") end - -@safetestset "DataLoader for input and output " begin include("data_loader/data_loader_for_input_and_output.jl") end - -@safetestset "HDF5 save/load for GML special array types " begin include("hdf5_support.jl") end +@safetestset "Matching " begin + include("data/test_matching.jl") +end + +@safetestset "Test data loader for q and p data " begin + include("data_loader/batch_data_loader_qp_test.jl") +end +@safetestset "Test mnist_utils. " begin + include("data_loader/mnist_utils.jl") +end +@safetestset "Test the data loader in combination with optimization_step! " begin + include("data_loader/data_loader_optimization_step.jl") +end +@safetestset "Optimizer functor with data loader for Adam " begin + include("data_loader/optimizer_functor_with_adam.jl") +end +@safetestset "Test data loader for a tensor (q and p data) " begin + include("data_loader/draw_batch_for_tensor_test.jl") +end + +@safetestset "Test NetworkLoss + Optimizer " begin + include("network_losses/losses_and_optimization.jl") +end + +@safetestset "Test parallel inverses " begin + include("kernels/tensor_inverse.jl") +end +@safetestset "Test parallel Cayley " begin + include("kernels/tensor_cayley.jl") +end + +@safetestset "Test volume-preserving feedforward neural network " begin + include("layers/volume_preserving_feedforward.jl") +end + +@safetestset "SympNet integrator " begin + include("sympnet_integrator.jl") +end +@safetestset "Regular transformer integrator " begin + include("standard_transformer_integrator.jl") +end + +@safetestset "Batch functor(s) " begin + include("batch/batch_functor.jl") +end + +@safetestset "Volume-Preserving Transformer (skew-symmetric tests) " begin + include("volume_preserving_attention/test_skew_map.jl") +end +@safetestset "Volume-Preserving Transformer (cayley-transform tests) " begin + include("volume_preserving_attention/test_cayley_transforms.jl") +end + +@safetestset "Linear Symplectic Attention " begin + include("linear_symplectic_attention.jl") +end +@safetestset "Linear Symplectic Transformer " begin + include("linear_symplectic_transformer.jl") +end + +@safetestset "DataLoader for input and output " begin + include("data_loader/data_loader_for_input_and_output.jl") +end + +@safetestset "HDF5 save/load for GML special array types " begin + include("hdf5_support.jl") +end diff --git a/test/sae_error_lower_than_psd_error.jl b/test/sae_error_lower_than_psd_error.jl index 2bd56a5a8..952c18a80 100644 --- a/test/sae_error_lower_than_psd_error.jl +++ b/test/sae_error_lower_than_psd_error.jl @@ -12,7 +12,7 @@ function test_accuracy(N::Integer, n::Integer; tol::Real = .35, n_epochs::Intege sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n; n_encoder_layers = 5, n_decoder_layers = 5)) - o = Optimizer(AdamOptimizer(), sae_nn) + o = Optimizer(Adam(), sae_nn) sae_error = o(sae_nn, dl, Batch(10), n_epochs)[end] @test sae_error < psd_error diff --git a/test/symplectic_autoencoder_tests.jl b/test/symplectic_autoencoder_tests.jl index 4c1f26aaa..2006ef0f9 100644 --- a/test/symplectic_autoencoder_tests.jl +++ b/test/symplectic_autoencoder_tests.jl @@ -1,23 +1,22 @@ using GeometricMachineLearning using Test using Zygote: jacobian -import Random +import Random Random.seed!(123) -function test_accuracy(N::Integer, n::Integer; tol::Real = .35, n_epochs::Integer = 100) +function test_accuracy(N::Integer, n::Integer; tol::Real = 0.35, n_epochs::Integer = 100) dl = DataLoader(rand(N, 10 * N); autoencoder = true) sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) - - o = Optimizer(AdamOptimizer(), sae_nn) + + o = Optimizer(Adam(), sae_nn) sae_error = o(sae_nn, dl, Batch(10), n_epochs)[end] - @test sae_error < tol + @test sae_error < tol end function test_encoder_and_decoder(N::Integer, n::Integer) - sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) sae_encoder = encoder(sae_nn) sae_decoder = decoder(sae_nn) @@ -33,14 +32,14 @@ function test_symplecticity(N::Integer, n::Integer) sae_nn = NeuralNetwork(SymplecticAutoencoder(N, n)) sae_decoder = decoder(sae_nn) test_vector = rand(n) - + # this matrix should be symplectic sympl_mat = jacobian(vec -> sae_decoder(vec), test_vector)[1] @test PoissonTensor(n) ≈ sympl_mat' * PoissonTensor(N) * sympl_mat # test if it's still symplectic after training dl = DataLoader(rand(N, 10 * N); autoencoder = true) - o = Optimizer(AdamOptimizer(), sae_nn) + o = Optimizer(Adam(), sae_nn) o(sae_nn, dl, Batch(10), 10) sympl_mat = jacobian(vec -> sae_decoder(vec), test_vector)[1] @test PoissonTensor(n) ≈ sympl_mat' * PoissonTensor(N) * sympl_mat @@ -80,4 +79,4 @@ all_tests(20, 10) # test the nonlinear decoder output layers test_decoder_output_layers(10, 6, 1) test_decoder_output_layers(10, 6, 2) -test_decoder_output_layers(20, 10, 3) \ No newline at end of file +test_decoder_output_layers(20, 10, 3) diff --git a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl index 216ba0c64..b0338cac2 100644 --- a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl +++ b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl @@ -1,59 +1,60 @@ -using GeometricMachineLearning, Test +using GeometricMachineLearning, GeometricOptimizers, Test import Random, LinearAlgebra Random.seed!(1234) @doc raw""" -This checks if the Adam cache was set up in the right way +This checks if the Adam cache was set up in the right way. +AdamCache fields (from GeometricOptimizers): x, g, δ, Δg, m₁, m₂, m̃₂, section. +The direction δ and moments m₁, m₂ should be zero-initialised. """ -function check_adam_cache(C::AbstractCache{T}, tol= T(10) * eps(T)) where T - @test typeof(C) <: AdamCache - @test propertynames(C) == (:B₁, :B₂) - @test typeof(C.B₁) <: StiefelLieAlgHorMatrix - @test typeof(C.B₂) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.B₁) < tol - @test LinearAlgebra.norm(C.B₂) < tol -end +_check_slahm_zero(A::StiefelLieAlgHorMatrix{T}, tol) where T = + (@test typeof(A) <: StiefelLieAlgHorMatrix; @test LinearAlgebra.norm(A) < tol) +_check_slahm_zero(A::NamedTuple, tol) = foreach(v -> _check_slahm_zero(v, tol), values(A)) + +function check_adam_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T + @test C isa AdamCache + _check_slahm_zero(C.δ, tol) + _check_slahm_zero(C.m₁, tol) +end check_adam_cache(B::NamedTuple) = apply_toNT(check_adam_cache, B) @doc raw""" -This checks if the momentum cache was set up in the right way +This checks if the momentum cache was set up in the right way. +MomentumCache fields (from GeometricOptimizers): x, g, δ, Δg, section. +The direction δ should be zero-initialised. """ -function check_momentum_cache(C::AbstractCache{T}, tol= T(10) * eps(T)) where T - @test typeof(C) <: MomentumCache - @test propertynames(C) == (:B,) - @test typeof(C.B) <: StiefelLieAlgHorMatrix - @test LinearAlgebra.norm(C.B) < tol +function check_momentum_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * eps(T)) where T + @test C isa MomentumCache + _check_slahm_zero(C.δ, tol) end check_momentum_cache(B::NamedTuple) = apply_toNT(check_momentum_cache, B) @doc raw""" -This checks if the gradient cache was set up in the right way +This checks if the gradient cache was set up in the right way. +GradientCache fields (from GeometricOptimizers): x, g, δ, Δg, section. """ -function check_gradient_cache(C::AbstractCache{T}) where T - @test typeof(C) <: GradientCache - @test propertynames(C) == () +function check_gradient_cache(C::GeometricOptimizers.OptimizerCache{T}) where T + @test C isa GradientCache + @test hasproperty(C, :δ) end check_gradient_cache(B::NamedTuple) = apply_toNT(check_gradient_cache, B) @doc raw""" This checks if all the caches are set up in the right way for the `MultiHeadAttention` layer with Stiefel weights. - -TODO: -- [ ] `BFGSOptimizer` !! """ function test_cache_setups_for_optimizer_for_multihead_attention_layer(T::Type, dim::Int, n_heads::Int) @assert dim % n_heads == 0 model = Chain(MultiHeadAttention(dim, n_heads, Stiefel=true)) ps = NeuralNetwork(model, CPU(), T).params - o₁ = Optimizer(AdamOptimizer(), ps) - o₂ = Optimizer(MomentumOptimizer(), ps) - o₃ = Optimizer(GradientOptimizer(), ps) + o₁ = Optimizer(Adam(), ps) + o₂ = Optimizer(MomentumMethod(), ps) + o₃ = Optimizer(GradientMethod(), ps) check_adam_cache(o₁.cache) check_momentum_cache(o₂.cache) check_gradient_cache(o₃.cache) end -test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) \ No newline at end of file +test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) diff --git a/test/transformer_related/multi_head_attention_stiefel_retraction.jl b/test/transformer_related/multi_head_attention_stiefel_retraction.jl index 400ab6478..6c2840213 100644 --- a/test/transformer_related/multi_head_attention_stiefel_retraction.jl +++ b/test/transformer_related/multi_head_attention_stiefel_retraction.jl @@ -1,7 +1,6 @@ -using GeometricMachineLearning, Test +using GeometricMachineLearning, GeometricOptimizers, Test using GeometricMachineLearning: geodesic using GeometricMachineLearning: cayley -using GeometricMachineLearning: init_optimizer_cache import Random, Test, LinearAlgebra, KernelAbstractions Random.seed!(1234) @@ -15,7 +14,7 @@ function check_retraction_geodesic(A::AbstractMatrix{T}, tol=eps(T)) where T @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end check_retraction_geodesic(cache::NamedTuple) = apply_toNT(check_retraction_geodesic, cache) -check_retraction_geodesic(B::MomentumCache) = check_retraction_geodesic(B.B) +check_retraction_geodesic(B::MomentumCache) = check_retraction_geodesic(B.δ) @doc raw""" This function computes the cayley retraction of an element of `StiefelLieAlgHorMatrix` and then checks if the resulting element is `StiefelProjection`. @@ -26,7 +25,7 @@ function check_retraction_cayley(A::AbstractMatrix{T}, tol=eps(T)) where T @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end check_retraction_cayley(cache::NamedTuple) = apply_toNT(check_retraction_cayley, cache) -check_retraction_cayley(B::MomentumCache) = check_retraction_cayley(B.B) +check_retraction_cayley(B::MomentumCache) = check_retraction_cayley(B.δ) @doc raw""" This is a test for that checks if the retractions (geodesic and Cayley for now) map from `StiefelLieAlgHorMatrix` to `StiefelManifold` when used with `MultiHeadAttention`. @@ -35,7 +34,7 @@ function test_multi_head_attention_retraction(T::Type, dim, n_heads, tol=eps(T), model = Chain(MultiHeadAttention(dim, n_heads, Stiefel=true)) ps = NeuralNetwork(model, backend, T).params - cache = init_optimizer_cache(MomentumOptimizer(), ps) + cache = Optimizer(MomentumMethod(), ps).cache check_retraction_geodesic(cache) diff --git a/test/transformer_related/multi_head_attention_stiefel_setup.jl b/test/transformer_related/multi_head_attention_stiefel_setup.jl index 67cf438a8..387c8675a 100644 --- a/test/transformer_related/multi_head_attention_stiefel_setup.jl +++ b/test/transformer_related/multi_head_attention_stiefel_setup.jl @@ -1,5 +1,4 @@ -using GeometricMachineLearning, Test -using GeometricMachineLearning: init_optimizer_cache +using GeometricMachineLearning, GeometricOptimizers, Test import Random, Test, LinearAlgebra, KernelAbstractions Random.seed!(1234) @@ -22,7 +21,7 @@ function check_grad_setup(B::AbstractMatrix{T}, tol=T(10)*eps(T)) where T @test LinearAlgebra.norm(B) < tol end check_grad_setup(gx::NamedTuple) = apply_toNT(check_grad_setup, gx) -check_grad_setup(B::MomentumCache) = check_grad_setup(B.B) +check_grad_setup(B::MomentumCache) = check_grad_setup(B.δ) @doc raw""" Check if `initialparameters` and `init_optimizer_cache` do the right thing for `MultiHeadAttentionLayer`. @@ -33,7 +32,7 @@ function check_multi_head_attention_stiefel_setup(T::Type, N::Int, n::Int) check_setup(ps) - gx = init_optimizer_cache(MomentumOptimizer(), ps) + gx = Optimizer(MomentumMethod(), ps).cache check_grad_setup(gx) end diff --git a/test/transformer_related/transformer_optimizer.jl b/test/transformer_related/transformer_optimizer.jl index d6907d577..b612e6847 100644 --- a/test/transformer_related/transformer_optimizer.jl +++ b/test/transformer_related/transformer_optimizer.jl @@ -5,39 +5,36 @@ import Random Random.seed!(1234) @doc raw""" -This function tests if the `GradientOptimzier`, `MomentumOptimizer`, `AdamOptimizer` and `BFGSOptimizer` act on the neural network weights via `optimization_step!`. +This function tests if the `GradientMethod`, `MomentumMethod`, and `Adam` act on the neural network weights via `optimization_step!`. """ function transformer_gradient_test(T, dim, n_heads, L, seq_length=8, batch_size=10) model = Chain(Transformer(dim, n_heads, L, Stiefel=true), ResNetLayer(dim)) model = Transformer(dim, n_heads, L, Stiefel=true) ps = NeuralNetwork(model, KernelAbstractions.CPU(), T).params - + input = rand(T, dim, seq_length, batch_size) - + loss(ps, input) = norm(model(input, ps)) dx = Zygote.gradient(ps -> loss(ps, input), ps)[1] - o₁ = Optimizer(GradientOptimizer(), ps) - o₂ = Optimizer(MomentumOptimizer(), ps) - o₃ = Optimizer(AdamOptimizer(), ps) - o₄ = Optimizer(BFGSOptimizer(), ps) + o₁ = Optimizer(GradientMethod(), ps) + o₂ = Optimizer(MomentumMethod(), ps) + o₃ = Optimizer(Adam(), ps) + # BFGSOptimizer is not available in GeometricOptimizers manifold optimizers ps₁ = deepcopy(ps) ps₂ = deepcopy(ps) ps₃ = deepcopy(ps) - ps₄ = deepcopy(ps) λY₁ = GlobalSection(ps₁) λY₂ = GlobalSection(ps₂) λY₃ = GlobalSection(ps₃) - λY₄ = GlobalSection(ps₄) optimization_step!(o₁, λY₁, ps₁, dx) optimization_step!(o₂, λY₂, ps₂, dx) optimization_step!(o₃, λY₃, ps₃, dx) - optimization_step!(o₄, λY₄, ps₄, dx) - @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps₄) == typeof(ps) - @test ps₁[1].PQ.head_1 ≉ ps₂[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 ≉ ps₄[1].PQ.head_1 # ≉ ps[1].PQ.head_1 + @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps) + @test ps₁[1].PQ.head_1 ≉ ps₂[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 # ≉ ps[1].PQ.head_1 end transformer_gradient_test(Float32, 10, 5, 4) \ No newline at end of file