NeuralNetworks in GeometricMachineLearning

GeometricMachineLearning inherits some functionality from another Julia package called AbstractNeuralNetworks. How these two packages interact is shown in the figure below for the example of the SympNet[1]:

Visualization of how the packages interact. Visualization of how the packages interact.

The red color indicates an abstract type, blue indicates a struct and orange indicates a const (derived from a struct). Solid black arrows indicate direct dependencies, i.e. we have

GradientLayer <: SympNetLayer <: AbstractExplicitLayer
GSympNet <: SympNet <: Architecture

Dashed black arrows indicate a derived neural network architecture. A GSympNet (which is an Architecture) is derived from GradientLayerQ and GradientLayerP (which are AbstractExplicitLayers) for example. An Architecture can be turned into a NeuralNetwork by calling the associated constructor

arch = GSympNet(3)
nn = NeuralNetwork(arch, CPU(), Float64)

Such a neural network has four fields:

  • architecture: the Architecture we supplied the constructor with,
  • model: a translation of the supplied architecture into specific neural network layers,
  • params: the neural network parameters,
  • backend: this indicates on which device we allocate the neural network parameters. In this case it is CPU().

We can get the associated model to GSympNet by calling:

nn.model.layers
(GradientLayerQ{3, 3, typeof(tanh)}(6, tanh), GradientLayerP{3, 3, typeof(tanh)}(6, tanh))

and we see that it consists of two layers: a GradientLayerQ and a GradientLayerP.

Saving and Loading

GeometricMachineLearning adds HDF5-backed save and load methods for NeuralNetwork to the generics NeuralNetworkParameters defines. Writing and reading the parameter set itself belongs to that package; the structured parameter types (StiefelManifold, SymmetricMatrix, SkewSymMatrix, …) come back as themselves because GeometricOptimizers, which owns them, registers how each is rebuilt. Passing a prototype parameter set to load rebuilds against it and needs no registration at all.

NeuralNetworkParameters.saveMethod
save(h5::HDF5.H5DataStore, nn::NeuralNetwork)

Save the parameters of nn into an already-open HDF5 store.

Extends save with a dispatch on NeuralNetwork. The parameters themselves are written by NeuralNetworkParameters, which tags each structured leaf with the type to rebuild it as and records the key order of every group.

source
NeuralNetworkParameters.loadMethod
load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; backend = CPU())
load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; backend = CPU())

Load network parameters from an already-open HDF5 store and return a NeuralNetwork for arch.

The element type is whatever the file holds, so a Float32 network reloads as Float32.

Structured parameters — StiefelManifold, SymmetricMatrix and the rest — are rebuilt from the type each was stored under, which GeometricOptimizers registers with NeuralNetworkParameters.register_parameter_type!. Pass prototype, a parameter set of the right shape, to rebuild against it instead and skip the registry altogether.

source
NeuralNetworkParameters.loadMethod
load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; backend = CPU())
load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; backend = CPU())

Convenience overload: open filename for reading, then call load on the store.

source
  • 1The section on SympNets also contains an explanation of all the structs and types described in this section here.