AbstractNeuralNetworks

NeuralNetworkBackend

The backend that specifies where and how neural network parameters are allocated.

It largely inherits properties from KernelAbstractions.Backend, but also adds CPUStatic which is defined in AbstractNeuralNetworks.

AbstractCell

An AbstractCell is a map from $\mathbb{R}^{M}×\mathbb{R}^{N} \rightarrow \mathbb{R}^{O}×\mathbb{R}^{P}$.

Concrete cell types should implement the following functions:

• initialparameters(backend::NeuralNetworkBackend, ::Type{T}, cell::AbstractCell; init::Initializer = default_initializer(), rng::AbstractRNG = Random.default_rng())
• update!(::AbstractLayer, θ::NamedTuple, dθ::NamedTuple, η::AbstractFloat)

and the functors

• cell(x, st, ps)
• cell(z, y, x, st, ps)

AbstractExplicitCell

Abstract supertype for explicit cells. This type exists mainly for compatibility with Lux.

AbstractExplicitLayer

Abstract supertype for explicit layers. This type exists mainly for compatibility with Lux.

AbstractLayer

An AbstractLayer is a map from $\mathbb{R}^{M} \rightarrow \mathbb{R}^{N}$.

Concrete layer types should implement the following functions:

• initialparameters(backend::NeuralNetworkBackend, ::Type{T}, layer::AbstractLayer; init::Initializer = default_initializer(), rng::AbstractRNG = Random.default_rng())
• update!(::AbstractLayer, θ::NamedTuple, dθ::NamedTuple, η::AbstractFloat)

and the functors

• layer(x, ps)
• layer(y, x, ps)

AbstractPullback{NNLT<:NetworkLoss}

AbstractPullback is an abstract type that encompasses all ways of performing differentiation (especially computing the gradient with respect to neural network parameters) in GeometricMachineLearning.

If a user wants to implement a custom Pullback the following two functions have to be extended:

(_pullback::AbstractPullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})
(_pullback::AbstractPullback)(ps, model, input_nt::QPT)

based on the loss::NetworkLoss that's stored in _pullback. The output of _pullback needs to be a Tuple that contains:

1. the loss evaluated at ps and input_nt (or input_nt_output_nt),
2. the gradient of loss with respect to ps that call be called with e.g.:

_pullback(ps, model, input_nt)[2](1) # returns the gradient wrt to `ps`

$\ldots$ we use this convention as it is analogous to how Zygote builds pullbacks.

An example is GeometricMachineLearning.ZygotePullback.

Architecture

CPUStatic

An additional backend that specifies allocation of static arrays.

Implementation

This is not a subtype of KernelAbstractions.Backend as it is associated with StaticArrays.MArray and such subtyping would therefore constitute type piracy.

Chain

A chain is a sequence of layers.

A Chain can be initialized by passing an arbitrary number of layers

Chain(layers...)

or a neural network architecture together with a backend and a parameter type:

Chain(::Architecture, ::NeuralNetworkBackend, ::Type; kwargs...)
Chain(::Architecture, ::Type; kwargs...)

If the backend is omitted, the default backend CPU() is chosen. The keyword arguments will be passed to the initialparameters method of each layer.

FeedForwardLoss()

Make an instance of a loss for feedforward neural networks.

This should be used together with a neural network of type GeometricMachineLearning.NeuralNetworkIntegrator.

Example

FeedForwardLoss applies a neural network to an input and compares it to the output via an $L_2$ norm:

using AbstractNeuralNetworks
using LinearAlgebra: norm
import Random
Random.seed!(123)

const d = 2
arch = Chain(Dense(d, d), Dense(d, d))
nn = NeuralNetwork(arch)

input_vec =  [1., 2.]
output_vec = [3., 4.]
loss = FeedForwardLoss()

loss(nn, input_vec, output_vec) ≈ norm(output_vec - nn(input_vec)) / norm(output_vec)

# output

true

So FeedForwardLoss simply does:

\[ \mathtt{loss}(\mathcal{NN}, \mathtt{input}, \mathtt{output}) = || \mathcal{NN}(\mathtt{input}) - \mathtt{output} || / || \mathtt{output}||,\]

where $||\cdot||$ is the $L_2$ norm.

Parameters

This loss does not have any parameters.

GlorotUniform <: Initializer

Glorot uniform was introduced by [1].

Initializer

Determines how neural network weights are initialized.

A supertype for Chain, AbstractCell etc.

NetworkLoss

An abstract type for all the neural network losses. If you want to implement CustomLoss <: NetworkLoss you need to define a functor:

(loss::CustomLoss)(model, ps, input, output)

where model is an instance of an AbstractExplicitLayer or a Chain and ps the parameters.

See FeedForwardLoss, GeometricMachineLearning.TransformerLoss, GeometricMachineLearning.AutoEncoderLoss and GeometricMachineLearning.ReducedLoss for examples.

NeuralNetwork <: AbstractNeuralNetwork

Neuralnetwork stores the Architecture, Model, neural network paramters and backend of the system.

Implementation

See NeuralNetworkBackend for the backend.

NeuralNetworkParameters

This struct stores the parameters of a neural network. In essence, it is just a wrapper around a NamedTuple of NamedTuple that provides some context, e.g., for storing parameters to file.

OneInitializer <: Initializer

ZeroInitializer <: Initializer

apply!(y, cell::AbstractCell, x, ps)

Simply calls cell(y, x, ps)

apply!(y, layer::AbstractLayer, x, ps)

Simply calls layer(y, x, ps)

apply(cayer::AbstractCell, x, ps)

Simply calls cell(x, st, ps)

apply(layer::AbstractLayer, x, ps)

Simply calls layer(x, ps)

changebackend(backend, nn)

Extended help

The function changebackend is defined for NeuralNetworkParameters, NeuralNetwork, AbstractArrays and NamedTuples. This function is also exported.

initialparameters

Returns the initial parameters of a model, i.e., a layer or chain.

initialparameters(backend::NeuralNetworkBackend, ::Type{T}, model::Model; init::Initializer = default_initializer(), rng::AbstractRNG = Random.default_rng())
initialparameters(::Type{T}, model::Model; init::Initializer = default_initializer(), rng::AbstractRNG = Random.default_rng())

The init! function must have the following signature:

init!(rng::AbstractRNG, x::AbstractArray)

The default_initializer() returns randn!.

networkbackend(arr)

Returns the NeuralNetworkBackend of arr.