SimpleChains.jl

https://pumasai.github.io/SimpleChains.jl/stable
Admin
SimpleChains.jl is a Julia package for building and training neural networks, offering various...
Tokens:7,621
Snippets:61
Trust Score:8.3
Update:11 months ago
Show doc for...
Tokens:
Raw
### Get Parameters and Biases from SimpleChain

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Retrieves views of the parameters and biases from a SimpleChain 'sc' within a given parameter vector 'p'. The 'inputdim' argument is optional.

```julia
params(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
biases(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
```

--------------------------------

### SimpleChains Flatten Layer Example

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Flattens the first N dimensions of an input tensor. The example demonstrates flattening the first two dimensions of a 3D array.

```julia
SimpleChains.Flatten{N}()
```

```julia
julia> Flatten{2}()(rand(2, 3, 4))
6×4 Matrix{Float64}:
 0.0609115  0.597285  0.279899  0.888223
 0.0667422  0.315741  0.351003  0.805629
 0.678297   0.350817  0.984215  0.399418
 0.125801   0.566696  0.96873   0.57744
 0.331961   0.350742  0.59598   0.741998
 0.26345    0.144635  0.076433  0.330475
```

--------------------------------

### SimpleChains AbstractPenalty Interface

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Defines the interface for penalty functions, which are used for regularization. Implementations must provide methods to get the associated chain, apply the penalty, and optionally update gradients.

```julia
SimpleChains.AbstractPenalty

# Methods:
getchain(::AbstractPenalty)::SimpleChain
apply_penalty(::AbstractPenalty, params)::Number
apply_penalty!(grad, ::AbstractPenalty, params)::Number
```

--------------------------------

### Get Front of SimpleChain

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Returns the SimpleChain 'c' without its last layer, useful for removing a loss layer.

```julia
Base.front(c::SimpleChain)
```

--------------------------------

### Get Layer Output Dimension

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Calculates the number of parameters required by a layer ('d') given an input dimension 'inputdim'. It also returns the output size of the layer, which serves as the input dimension for the subsequent layer.

```julia
numparam(d::Layer, inputdim::Tuple)
```

--------------------------------

### SimpleChains.init_params

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Creates and initializes a parameter vector for a SimpleChain.

```APIDOC
## SimpleChains.init_params

### Description
Creates a parameter vector of element type `T` with size matching that by `id` (argument not required if provided to the `chain` object itself). See the documentation of the individual layers to see how they are initialized, but it is generally via (Xavier) Glorot uniform or normal distributions.

### Method
`SimpleChains.init_params(chn[, id = nothing][, ::Type{T} = Float32])`

### Endpoint
N/A (Function)
```

--------------------------------

### SimpleChains Parameter Initialization and Training Utilities

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Provides utility functions for initializing parameters of a SimpleChain, calculating gradients, and training the model using batched or unbatched data.

```julia
SimpleChains.init_params(chain)
SimpleChains.init_params!(chain)
SimpleChains.train_batched!(...) 
SimpleChains.train_unbatched!(...)
SimpleChains.valgrad!(...)
```

--------------------------------

### Initialize Parameters, Loss Function, and Training Loop

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/smallmlp

Initializes the parameters for the MLP, allocates memory for gradients, defines the squared loss function for both training and testing sets, and sets up a reporting function to monitor loss. It then enters a training loop using the ADAM optimizer.

```julia
@time p = SimpleChains.init_params(mlpd);
G = SimpleChains.alloc_threaded_grad(mlpd);

mlpdloss = SimpleChains.add_loss(mlpd, SquaredLoss(Y));
mlpdtest = SimpleChains.add_loss(mlpd, SquaredLoss(Ytest));

# define a function named report to calculate and report the value of loss function with train and test sets.
report = let mtrain = mlpdloss, X=X, Xtest=Xtest, mtest = mlpdtest
  p -> begin
    let train = mtrain(X, p), test = mtest(Xtest, p)
      @info "Loss:" train test
    end
  end
end

report(p)
for _ in 1:3
  @time SimpleChains.train_unbatched!(
    G, p, mlpdloss, X, SimpleChains.ADAM(), 10_000
  );
  report(p)
end
```

--------------------------------

### SimpleChains.init_params!

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Randomly initializes the parameter vector `p` of a SimpleChain, typically using Glorot distributions.

```APIDOC
## SimpleChains.init_params!

### Description
Randomly initializes parameter vector `p` with input dim `id`. Input dim does not need to be specified if these were provided to the chain object itself. See the documentation of the individual layers to see how they are initialized, but it is generally via (Xavier) Glorot uniform or normal distributions.

### Method
`SimpleChains.init_params!(chn, p, id = nothing)`

### Endpoint
N/A (Function)
```

--------------------------------

### Base.front

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Retrieves the front of a SimpleChain, which is useful for popping off a loss layer.

```APIDOC
## Base.front SimpleChain

### Description
Retrieves the front of a SimpleChain, which is useful for popping off a loss layer.

### Method
`Base.front(c::SimpleChain)`

### Endpoint
N/A (Function)
```

--------------------------------

### Initialize Parameters for SimpleChain

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Initializes the parameter vector 'p' for a SimpleChain 'chn'. If 'id' is not provided, it uses dimensions from the chain object. Initialization typically uses Xavier Glorot distributions for weights and zero for biases.

```julia
SimpleChains.init_params!(chn, p, id = nothing)
SimpleChains.init_params(chn[, id = nothing][, ::Type{T} = Float32])
```

--------------------------------

### SimpleChains.train_batched!

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Trains a SimpleChain using batched input data, updating parameters in place.

```APIDOC
## SimpleChains.train_batched!

### Description
Train while batching arguments.

### Method
`train_batched!(g::AbstractVecOrMat, p, chn, X, opt, iters; batchsize = nothing)`

### Parameters
#### Arguments
* `g`: Pre-allocated gradient buffer. Can be allocated with `similar(p)` (if you want to run single threaded), or `alloc_threaded_grad(chn, size(X))` (`size(X)` argument is only necessary if the input dimension was not specified when constructing the chain). If a matrix, the number of columns gives how many threads to use. Do not use more threads than batch size would allow.
* `p`: The parameter vector. It is updated inplace. It should be pre-initialized, e.g. with `init_params`/`init_params!`.
* `chn`: The `SimpleChain`. It must include a loss (see `SimpleChains.add_loss`) containing the target information (dependent variables) you're trying to fit.
* `X`: The training data input argument (independent variables).
* `opt`: The optimizer. Currently, only `SimpleChains.ADAM` is supported.
* `iters`: How many iterations to train for.
* `batchsize` (keyword argument): The size of the batches to use. If `batchsize = nothing`, it'll try to do a half-decent job of picking the batch size for you. However, this is not well optimized at the moment.

### Endpoint
N/A (Function)
```

--------------------------------

### Use Custom Loss Function in Model Training

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/custom_loss_layer

Demonstrates how to use the custom BinaryLogitCrossEntropyLoss in a SimpleChains model. Includes model creation, data generation, parameter initialization, and training using valgrad! and train_unbatched! functions.

```julia
using SimpleChains

model = SimpleChain(
    static(2),
    TurboDense(tanh, 32),
    TurboDense(tanh, 16),
    TurboDense(identity, 1)
)

batch_size = 64
X = rand(Float32, 2, batch_size)
Y = rand(Bool, batch_size)

parameters = SimpleChains.init_params(model);
gradients = SimpleChains.alloc_threaded_grad(model);

# Add the loss like any other loss type
model_loss = SimpleChains.add_loss(model, BinaryLogitCrossEntropyLoss(Y));

SimpleChains.valgrad!(gradients, model_loss, X, parameters)

epochs = 100
SimpleChains.train_unbatched!(gradients, parameters, model_loss, X, SimpleChains.ADAM(), epochs);
```

--------------------------------

### SimpleChains.train_unbatched!

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Trains a SimpleChain using unbatched input data, updating parameters in place.

```APIDOC
## SimpleChains.train_unbatched!

### Description
Train without batching inputs.

### Method
`train_unbatched!([g::AbstractVecOrMat, ]p, chn, X, opt, iters)`

### Parameters
#### Arguments
* `g` (optional): Pre-allocated gradient buffer. Can be allocated with `similar(p)` (if you want to run single threaded), or `alloc_threaded_grad(chn, size(X))` (`size(X)` argument is only necessary if the input dimension was not specified when constructing the chain). If a matrix, the number of columns gives how many threads to use. Do not use more threads than batch size would allow. This argument is optional. If excluded, it will run multithreaded (assuming you started Julia with multiple threads).
* `p`: The parameter vector. It is updated inplace. It should be pre-initialized, e.g. with `init_params`/`init_params!`.
* `chn`: The `SimpleChain`. It must include a loss (see `SimpleChains.add_loss`) containing the target information (dependent variables) you're trying to fit.
* `X`: The training data input argument (independent variables).
* `opt`: The optimizer. Currently, only `SimpleChains.ADAM` is supported.
* `iters`: How many iterations to train for.

### Endpoint
N/A (Function)
```

--------------------------------

### SimpleChains.params

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Extracts the parameters from a SimpleChain's parameter vector.

```APIDOC
## SimpleChains.params

### Description
Returns a tuple of the parameters of the SimpleChain `sc`, as a view of the parameter vector `p`.

### Method
`params(sc::SimpleChain, p::AbstractVector, inputdim = nothing)`

### Endpoint
N/A (Function)
```

--------------------------------

### Train SimpleChain with Batched Data

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Trains a SimpleChain model using batched input data. Requires a pre-allocated gradient buffer 'g', parameter vector 'p', the chain 'chn', training data 'X', an optimizer 'opt', and the number of iterations 'iters'. Batch size can be automatically determined or specified.

```julia
train_batched!(g::AbstractVecOrMat, p, chn, X, opt, iters; batchsize = nothing)
```

--------------------------------

### Initialize LeNet5 parameters with SimpleChains.jl

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/mnist

Creates a random parameter vector for the LeNet5 model using Glorot initialization. Can infer parameter dimensions from statically sized model or explicitly specify input dimensions. Returns a vector suitable for optimization algorithms.

```julia
@time p = SimpleChains.init_params(lenet);

# Alternative with explicit input size
@time p = SimpleChains.init_params(lenet, size(xtrain4));
```

--------------------------------

### SimpleChains SimpleChain Construction and Usage

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Constructs a SimpleChain, which is a sequence of layers for neural networks. It can optionally specify input dimensions and can be called with data and parameters to perform a forward pass.

```julia
SimpleChains.SimpleChain([inputdim::Union{Integer,Tuple{Vararg{Integer}}}, ] layers)
```

```julia
c = SimpleChain(...);
p = SimpleChains.init_params(c);
c(X, p) # X are the independent variables, and `p` the parameter vector.
```

--------------------------------

### Preallocate Gradient Buffer for Training

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Allocates a pre-initialized array for storing gradients, optimized for use with `train_batched` and `train_unbatched`. Supports multi-threading by creating a matrix with columns per thread for parallel gradient accumulation. Memory is aligned to prevent false sharing.

```julia
alloc_threaded_grad(chn, id = nothing, ::Type{T} = Float32; numthreads = min(Threads.nthreads(), SimpleChains.num_cores()))
```

--------------------------------

### Train SimpleChain without Batching

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Trains a SimpleChain model using unbatched input data. Accepts an optional pre-allocated gradient buffer 'g', parameter vector 'p', the chain 'chn', training data 'X', an optimizer 'opt', and the number of iterations 'iters'. If 'g' is omitted, multi-threading is enabled.

```julia
train_unbatched!([g::AbstractVecOrMat, ]p, chn, X, opt, iters)
```

--------------------------------

### SimpleChains FrontLastPenalty

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Applies different penalty functions to the front and last layers of a SimpleChain. It uses a 'frontpen' for all layers except the last, and a 'lastpen' for the final layer.

```julia
SimpleChains.FrontLastPenalty(SimpleChain, frontpen(λ₁...), lastpen(λ₂...))
```

--------------------------------

### SimpleChains Dropout Layer

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Implements a dropout layer for regularization. During evaluation without gradients, it scales inputs. With gradients, it randomly zeros a proportion of inputs based on the dropout rate 'p'.

```julia
SimpleChains.Dropout(p) # 0 < p < 1
```

--------------------------------

### SimpleChains.numparam

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Calculates the number of parameters and output dimension for a given layer and input dimension.

```APIDOC
## SimpleChains.numparam

### Description
Returns a `Tuple{Int,S}`. The first element is the number of parameters required by the layer given an argument of size `inputdim`. The second argument is the size of the object returned by the layer, which can be fed into `numparam` of the following layer.

### Method
`numparam(d::Layer, inputdim::Tuple)`

### Endpoint
N/A (Function)
```

--------------------------------

### SimpleChains Convolutional Layer (Conv)

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Performs a convolution operation with specified dimensions and output dimension, followed by bias addition and activation. Weights are initialized using Glorot uniform distribution and biases are zero-initialized.

```julia
SimpleChains.Conv(activation, dims::Tuple{Vararg{Integer}}, outputdim::Integer)
```

--------------------------------

### SimpleChains.biases

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Extracts the biases from a SimpleChain's parameter vector.

```APIDOC
## SimpleChains.biases

### Description
Returns a tuple of the biases of the SimpleChain `sc`, as a view of the parameter vector `p`.

### Method
`biases(sc::SimpleChain, p::AbstractVector, inputdim = nothing)`

### Endpoint
N/A (Function)
```

--------------------------------

### SimpleChains ADAM Optimizer Configuration

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Configures the ADAM optimizer with specified learning rate (η) and momentum parameters (β). This is used for training neural networks by updating model weights.

```julia
SimpleChains.ADAM(η = 0.001, β = (0.9, 0.999))
```

--------------------------------

### Import SimpleChains and Define Custom Loss Type

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/custom_loss_layer

Imports SimpleChains and defines a custom loss type called BinaryLogitCrossEntropyLoss which subtypes SimpleChains.AbstractLoss. This struct holds target values used for computing the loss.

```julia
using SimpleChains

struct BinaryLogitCrossEntropyLoss{T,Y<:AbstractVector{T}} <: SimpleChains.AbstractLoss{T}
    targets::Y
end

SimpleChains.target(loss::BinaryLogitCrossEntropyLoss) = loss.targets
(loss::BinaryLogitCrossEntropyLoss)(x::AbstractArray) = BinaryLogitCrossEntropyLoss(x)
```

--------------------------------

### Compute Value and Gradient - Julia

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

The valgrad! function computes the value and gradients for a SimpleChain model. It updates gradients in-place for parameters and optionally for inputs, supporting ADAM optimizer. Inputs include gradient storage, chain, arguments, and parameters; outputs are computed values and gradients. Limitations: Only ADAM is supported, and gradients must match sizes.

```julia
valgrad!(g, c::SimpleChain, arg, params)
```

--------------------------------

### Load and preprocess MNIST data with MLDatasets.jl

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/mnist

Loads MNIST training and test datasets using MLDatasets.jl. Reshapes 3D image arrays to 4D format with channel dimension and converts labels to 1-based indexing for SimpleChains compatibility. Prepares data for convolutional neural network training.

```julia
using MLDatasets
xtrain3, ytrain0 = MLDatasets.MNIST.traindata(Float32);
xtest3, ytest0 = MLDatasets.MNIST.testdata(Float32);
size(xtest3)
# (28, 28, 60000)
extrema(ytrain0) # digits, 0,...,9
# (0, 9)

xtrain4 = reshape(xtrain3, 28, 28, 1, :);
xtest4 = reshape(xtest3, 28, 28, 1, :);
ytrain1 = UInt32.(ytrain0 .+ 1);
ytest1 = UInt32.(ytest0 .+ 1);
```

--------------------------------

### Calculate Loss Function Implementation

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/custom_loss_layer

Implements the calculation of binary logit cross-entropy loss given logits and target values. Iterates over targets to compute total loss using sigmoid probabilities.

```julia
function calculate_loss(loss::BinaryLogitCrossEntropyLoss, logits)
    y = loss.targets
    total_loss = zero(eltype(logits))
    for i in eachindex(y)
        p_i = inv(1 + exp(-logits[i]))
        y_i = y[i]
        total_loss -= y_i * log(p_i) + (1 - y_i) * (1 - log(p_i))
    end
    total_loss
end
function (loss::BinaryLogitCrossEntropyLoss)(previous_layer_output::AbstractArray{T}, p::Ptr, pu) where {T}
    total_loss = calculate_loss(loss, previous_layer_output)
    total_loss, p, pu
end
```

--------------------------------

### Backpropagate Gradient from Custom Loss Function

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/custom_loss_layer

Implements gradient backpropagation for the custom loss function. Computes total loss and updates the previous layer's output with gradients based on targets and logits.

```julia
function SimpleChains.chain_valgrad!(
    __,
    previous_layer_output::AbstractArray{T},
    layers::Tuple{BinaryLogitCrossEntropyLoss},
    _::Ptr,
    pu::Ptr{UInt8},
) where {T}
    loss = getfield(layers, 1)
    total_loss = calculate_loss(loss, previous_layer_output)
    y = loss.targets

    # Store the backpropagated gradient in the previous_layer_output array.
    for i in eachindex(y)
        sign_arg = 2 * y[i] - 1
        # Get the value of the last logit
        logit_i = previous_layer_output[i]
        previous_layer_output[i] = -(sign_arg * inv(1 + exp(sign_arg * logit_i)))
    end

    return total_loss, previous_layer_output, pu
end
```

--------------------------------

### SimpleChains.add_loss

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Adds a loss function to a SimpleChain. The loss function should contain the target data for training.

```APIDOC
## SimpleChains.add_loss

### Description
Adds the loss function `l` to the simple chain. The loss function should hold the target you're trying to fit.

### Method
`add_loss(chn, l::AbstractLoss)`

### Endpoint
N/A (Function)
```

--------------------------------

### SimpleChains.alloc_threaded_grad

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Allocates a pre-initialized array for gradient accumulation, suitable for parallel training with `train_batched` and `train_unbatched`.

```APIDOC
## SimpleChains.alloc_threaded_grad

### Description
Returns a preallocated array for writing gradients, for use with `train_batched` and `train_unbatched`. If Julia was started with multiple threads, returns a matrix with one column per thread, so they may accumulate gradients in parallel. Note that the memory is aligned to avoid false sharing.

### Method
`alloc_threaded_grad(chn, id = nothing, ::Type{T} = Float32; numthreads = min(Threads.nthreads(), SimpleChains.num_cores()))`

### Endpoint
N/A (Function)
```

--------------------------------

### Define MLP Network Architecture with SimpleChains.jl

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/smallmlp

Defines a simple multi-layer perceptron (MLP) network using SimpleChains.jl. It consists of an input layer of size 4, followed by three hidden layers with ReLU and identity activation functions, and a final output layer of size 4.

```julia
using SimpleChains

mlpd = SimpleChain(
  static(4),
  TurboDense(tanh, 32),
  TurboDense(tanh, 16),
  TurboDense(identity, 4)
)
```

--------------------------------

### SimpleChains L1 and L2 Penalty Functions

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Defines L1 and L2 regularization penalties. L1Penalty penalizes the absolute value of parameters, while L2Penalty penalizes the square of parameters.

```julia
SimpleChains.L1Penalty(λ)
```

```julia
SimpleChains.L2Penalty(λ)
```

--------------------------------

### Compute Pullback for SimpleChain Layer

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Calculates the pullback of a layer with respect to the input 'A' and the gradient 'C̄'. The result is stored either in 'dest' or in 'A' itself, depending on the method signature.

```julia
pullback_arg!(dest, layer, C̄, A, p, pu, pu2)
pullback_arg!(layer, C̄, A, p, pu, pu2)
```

--------------------------------

### SimpleChains TurboDense Layer

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Implements a linear (dense) layer, optionally including a bias term specified by the boolean parameter 'B'. This layer performs a matrix multiplication with the input and adds a bias if enabled.

```julia
SimpleChains.TurboDense{B=true}(activation, outputdim::Integer)
```

--------------------------------

### Define LeNet5 CNN architecture using SimpleChains.jl

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/mnist

Constructs a LeNet5 convolutional neural network with convolutional layers, max pooling, and dense layers. Uses static dimension annotations for input size validation and adds logit cross-entropy loss for classification. Supports 10-class MNIST digit recognition.

```julia
using SimpleChains

lenet = SimpleChain(
  (static(28), static(28), static(1)),
  SimpleChains.Conv(SimpleChains.relu, (5, 5), 6),
  SimpleChains.MaxPool(2, 2),
  SimpleChains.Conv(SimpleChains.relu, (5, 5), 16),
  SimpleChains.MaxPool(2, 2),
  Flatten(3),
  TurboDense(SimpleChains.relu, 120),
  TurboDense(SimpleChains.relu, 84),
  TurboDense(identity, 10),
)

lenetloss = SimpleChains.add_loss(lenet, LogitCrossEntropyLoss(ytrain1));
```

--------------------------------

### SimpleChains.pullback_arg!

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Computes the pullback of a layer with respect to its input, storing the result in a destination array or the input array itself.

```APIDOC
## SimpleChains.pullback_arg!

### Description
Computes the pullback of `layer` with respect to `A` and `C̄`, storing the result in `dest`.
Computes the pullback of `layer` with respect to `A` and `C̄`, storing the result in `A`.

### Method
`pullback_arg!(dest, layer, C̄, A, p, pu, pu2)`
`pullback_arg!(layer, C̄, A, p, pu, pu2)`

### Endpoint
N/A (Function)
```

--------------------------------

### Define Layer Output Size Functions for Loss

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/custom_loss_layer

Defines layer output size functions that indicate no temporary arrays are needed during forward pass computations. These methods delegate to SimpleChains' internal helper function.

```julia
function SimpleChains.layer_output_size(::Val{T}, sl::BinaryLogitCrossEntropyLoss, s::Tuple) where {T}
    SimpleChains._layer_output_size_no_temp(Val{T}(), sl, s)
end
function SimpleChains.forward_layer_output_size(::Val{T}, sl::BinaryLogitCrossEntropyLoss, s) where {T}
    SimpleChains._layer_output_size_no_temp(Val{T}(), sl, s)
end
```

--------------------------------

### Train LeNet5 CNN and evaluate accuracy/loss

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/mnist

Trains the LeNet5 model using batched ADAM optimizer with specified learning rate and epochs. Computes accuracy and loss on both training and test datasets. Supports incremental training for additional epochs with continued parameter updates.

```julia
@time SimpleChains.train_batched!(G, p, lenetloss, xtrain4, SimpleChains.ADAM(3e-4), 10);
SimpleChains.accuracy_and_loss(lenetloss, xtrain4, p)
SimpleChains.accuracy_and_loss(lenetloss, xtest4, ytest1, p)

# Train for additional epochs
@time SimpleChains.train_batched!(G, p, lenetloss, xtrain4, SimpleChains.ADAM(3e-4), 10);
SimpleChains.accuracy_and_loss(lenetloss, xtrain4, p)
SimpleChains.accuracy_and_loss(lenetloss, xtest4, ytest1, p)
```

--------------------------------

### SimpleChains Loss Function Definitions

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Defines various loss functions used to evaluate the difference between predicted and target values. These include AbsoluteLoss, LogitCrossEntropyLoss, and SquaredLoss.

```julia
SimpleChains.AbsoluteLoss
```

```julia
SimpleChains.LogitCrossEntropyLoss
```

```julia
SimpleChains.SquaredLoss(target)
```

--------------------------------

### SimpleChains Activation Layer

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Applies a specified activation function elementwise to the input. This is a common layer in neural networks for introducing non-linearity.

```julia
SimpleChains.Activation(activation)
```

--------------------------------

### Define Matrix Exponential Approximation Function

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/smallmlp

Defines a function `f(x)` that approximates the matrix exponential. It takes a flattened vector, reshapes it into a square matrix, calculates its exponential, and returns the result as a flattened vector. This function is used to generate training and testing data.

```julia
function f(x)
  N = Base.isqrt(length(x))
  A = reshape(view(x, 1:N*N), (N,N))
  expA = exp(A)
  vec(expA)
end

T = Float32;
X = randn(T, 2*2, 10_000);
Y = reduce(hcat, map(f, eachcol(X)));
Xtest = randn(T, 2*2, 10_000);
Ytest = reduce(hcat, map(f, eachcol(Xtest)));
```

--------------------------------

### Create multi-threaded gradient matrix for training

Source: https://pumasai.github.io/SimpleChains.jl/stable/examples/mnist

Allocates a gradient matrix with rows equal to parameter count and columns per thread for parallel training. Estimates physical core count based on system architecture. Enables efficient multi-threaded gradient computation during optimization.

```julia
estimated_num_cores = (Sys.CPU_THREADS ÷ ((Sys.ARCH === :x86_64) + 1));
G = SimpleChains.alloc_threaded_grad(lenetloss);
```

--------------------------------

### Add Loss Function to SimpleChain

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Appends a loss function 'l' to a SimpleChain 'chn'. The loss function should contain the target variables for fitting.

```julia
add_loss(chn, l::AbstractLoss)
```

--------------------------------

### Extract Weights from SimpleChain - Julia

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

The weights function retrieves the weights (parameters excluding biases) from a SimpleChain as a view of the parameter vector. It requires the chain and parameter vector, with optional input dimension. Returns a tuple of weight views. Suitable for model inspection; does not modify inputs.

```julia
weights(sc::SimpleChain, p::AbstractVector, inputdim = nothing)
```

--------------------------------

### SimpleChains MaxPool Layer

Source: https://pumasai.github.io/SimpleChains.jl/stable/index

Performs a max pooling operation over specified dimensions. This layer is used to downsample feature maps while retaining important information.

```julia
SimpleChains.MaxPool(dims::Tuple{Vararg{Integer}})
```
Tokens:
Raw