Training a Simple LSTM

In this tutorial we will go over using a recurrent neural network to classify clockwise and anticlockwise spirals. By the end of this tutorial you will be able to:

1. Create custom Lux models.

2. Become familiar with the Lux recurrent neural network API.

3. Training using Optimisers.jl and Zygote.jl.

Package Imports

Note: If you wish to use AutoZygote() for automatic differentiation, add Zygote to your project dependencies and include using Zygote.

using ADTypes, Lux, JLD2, MLUtils, Optimisers, Printf, Reactant, Random

Dataset

We will use MLUtils to generate 500 (noisy) clockwise and 500 (noisy) anticlockwise spirals. Using this data we will create a MLUtils.DataLoader. Our dataloader will give us sequences of size 2 × seq_len × batch_size and we need to predict a binary value whether the sequence is clockwise or anticlockwise.

function create_dataset(; dataset_size=1000, sequence_length=50)
    # Create the spirals
    data = [MLUtils.Datasets.make_spiral(sequence_length) for _ in 1:dataset_size]
    # Get the labels
    labels = vcat(repeat([0.0f0], dataset_size ÷ 2), repeat([1.0f0], dataset_size ÷ 2))
    clockwise_spirals = [
        reshape(d[1][:, 1:sequence_length], :, sequence_length, 1) for
        d in data[1:(dataset_size ÷ 2)]
    ]
    anticlockwise_spirals = [
        reshape(d[1][:, (sequence_length + 1):end], :, sequence_length, 1) for
        d in data[((dataset_size ÷ 2) + 1):end]
    ]
    x_data = Float32.(cat(clockwise_spirals..., anticlockwise_spirals...; dims=3))
    return x_data, labels
end

function get_dataloaders(; dataset_size=1000, sequence_length=50)
    x_data, labels = create_dataset(; dataset_size, sequence_length)
    # Split the dataset
    (x_train, y_train), (x_val, y_val) = splitobs((x_data, labels); at=0.8, shuffle=true)
    # Create DataLoaders
    return (
        # Use DataLoader to automatically minibatch and shuffle the data
        DataLoader(
            collect.((x_train, y_train)); batchsize=128, shuffle=true, partial=false
        ),
        # Don't shuffle the validation data
        DataLoader(collect.((x_val, y_val)); batchsize=128, shuffle=false, partial=false),
    )
end

Creating a Classifier

We will be extending the Lux.AbstractLuxContainerLayer type for our custom model since it will contain a LSTM block and a classifier head.

We pass the field names lstm_cell and classifier to the type to ensure that the parameters and states are automatically populated and we don't have to define Lux.initialparameters and Lux.initialstates.

To understand more about container layers, please look at Container Layer.

struct SpiralClassifier{L,C} <: AbstractLuxContainerLayer{(:lstm_cell, :classifier)}
    lstm_cell::L
    classifier::C
end

We won't define the model from scratch but rather use the Lux.LSTMCell and Lux.Dense.

function SpiralClassifier(in_dims, hidden_dims, out_dims)
    return SpiralClassifier(
        LSTMCell(in_dims => hidden_dims), Dense(hidden_dims => out_dims, sigmoid)
    )
end

We can use default Lux blocks – Recurrence(LSTMCell(in_dims => hidden_dims) – instead of defining the following. But let's still do it for the sake of it.

Now we need to define the behavior of the Classifier when it is invoked.

function (s::SpiralClassifier)(
    x::AbstractArray{T,3}, ps::NamedTuple, st::NamedTuple
) where {T}
    # First we will have to run the sequence through the LSTM Cell
    # The first call to LSTM Cell will create the initial hidden state
    # See that the parameters and states are automatically populated into a field called
    # `lstm_cell` We use `eachslice` to get the elements in the sequence without copying,
    # and `Iterators.peel` to split out the first element for LSTM initialization.
    x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2)))
    (y, carry), st_lstm = s.lstm_cell(x_init, ps.lstm_cell, st.lstm_cell)
    # Now that we have the hidden state and memory in `carry` we will pass the input and
    # `carry` jointly
    for x in x_rest
        (y, carry), st_lstm = s.lstm_cell((x, carry), ps.lstm_cell, st_lstm)
    end
    # After running through the sequence we will pass the output through the classifier
    y, st_classifier = s.classifier(y, ps.classifier, st.classifier)
    # Finally remember to create the updated state
    st = merge(st, (classifier=st_classifier, lstm_cell=st_lstm))
    return vec(y), st
end

Using the @compact API

We can also define the model using the Lux.@compact API, which is a more concise way of defining models. This macro automatically handles the boilerplate code for you and as such we recommend this way of defining custom layers

function SpiralClassifierCompact(in_dims, hidden_dims, out_dims)
    return @compact(;
        lstm_cell=LSTMCell(in_dims => hidden_dims),
        classifier=Dense(hidden_dims => out_dims, sigmoid)
    ) do x::AbstractArray{T,3} where {T}
        x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2)))
        y, carry = lstm_cell(x_init)
        for x in x_rest
            y, carry = lstm_cell((x, carry))
        end
        @return vec(classifier(y))
    end
end

Defining Accuracy, Loss and Optimiser

Now let's define the binary cross-entropy loss. Typically it is recommended to use logitbinarycrossentropy since it is more numerically stable, but for the sake of simplicity we will use binarycrossentropy.

const lossfn = BinaryCrossEntropyLoss()

function compute_loss(model, ps, st, (x, y))
    ŷ, st_ = model(x, ps, st)
    loss = lossfn(ŷ, y)
    return loss, st_, (; y_pred=ŷ)
end

matches(y_pred, y_true) = sum((y_pred .> 0.5f0) .== y_true)
accuracy(y_pred, y_true) = matches(y_pred, y_true) / length(y_pred)

Training the Model

function main(model_type)
    dev = reactant_device()
    cdev = cpu_device()

    # Get the dataloaders
    train_loader, val_loader = get_dataloaders() |> dev

    # Create the model
    model = model_type(2, 8, 1)
    ps, st = Lux.setup(Random.default_rng(), model) |> dev

    train_state = Training.TrainState(model, ps, st, Adam(0.01f0))
    model_compiled = if dev isa ReactantDevice
        @compile model(first(train_loader)[1], ps, Lux.testmode(st))
    else
        model
    end
    ad = dev isa ReactantDevice ? AutoReactant() : AutoZygote()

    for epoch in 1:25
        # Train the model
        total_loss = 0.0f0
        total_samples = 0
        for (x, y) in train_loader
            (_, loss, _, train_state) = Training.single_train_step!(
                ad, lossfn, (x, y), train_state
            )
            total_loss += loss * length(y)
            total_samples += length(y)
        end
        @printf("Epoch [%3d]: Loss %4.5f\n", epoch, total_loss / total_samples)

        # Validate the model
        total_acc = 0.0f0
        total_loss = 0.0f0
        total_samples = 0

        st_ = Lux.testmode(train_state.states)
        for (x, y) in val_loader
            ŷ, st_ = model_compiled(x, train_state.parameters, st_)
            ŷ, y = cdev(ŷ), cdev(y)
            total_acc += accuracy(ŷ, y) * length(y)
            total_loss += lossfn(ŷ, y) * length(y)
            total_samples += length(y)
        end

        @printf(
            "Validation:\tLoss %4.5f\tAccuracy %4.5f\n",
            total_loss / total_samples,
            total_acc / total_samples
        )
    end

    return (train_state.parameters, train_state.states) |> cdev
end

ps_trained, st_trained = main(SpiralClassifier)

┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
└ @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:18
┌ Warning: You are using code generated by an older version of ProtoBuf.jl, which was deprecated. Please regenerate your protobuf definitions with the current version of ProtoBuf.jl. The new version will allow for defining custom AbstractProtoDecoder variants. This warning is only printed once per session.
└ @ ProtoBuf.Codecs ~/.julia/packages/ProtoBuf/85REE/src/codec/decode.jl:194
Epoch [  1]: Loss 0.71583
Validation:	Loss 0.64015	Accuracy 0.70312
Epoch [  2]: Loss 0.58941
Validation:	Loss 0.52401	Accuracy 1.00000
Epoch [  3]: Loss 0.48436
Validation:	Loss 0.41896	Accuracy 1.00000
Epoch [  4]: Loss 0.37822
Validation:	Loss 0.31871	Accuracy 1.00000
Epoch [  5]: Loss 0.29578
Validation:	Loss 0.25745	Accuracy 1.00000
Epoch [  6]: Loss 0.23930
Validation:	Loss 0.20704	Accuracy 1.00000
Epoch [  7]: Loss 0.19130
Validation:	Loss 0.16395	Accuracy 1.00000
Epoch [  8]: Loss 0.15027
Validation:	Loss 0.12868	Accuracy 1.00000
Epoch [  9]: Loss 0.11825
Validation:	Loss 0.09970	Accuracy 1.00000
Epoch [ 10]: Loss 0.09159
Validation:	Loss 0.07773	Accuracy 1.00000
Epoch [ 11]: Loss 0.07272
Validation:	Loss 0.06244	Accuracy 1.00000
Epoch [ 12]: Loss 0.05927
Validation:	Loss 0.05194	Accuracy 1.00000
Epoch [ 13]: Loss 0.04980
Validation:	Loss 0.04436	Accuracy 1.00000
Epoch [ 14]: Loss 0.04313
Validation:	Loss 0.03862	Accuracy 1.00000
Epoch [ 15]: Loss 0.03798
Validation:	Loss 0.03417	Accuracy 1.00000
Epoch [ 16]: Loss 0.03370
Validation:	Loss 0.03054	Accuracy 1.00000
Epoch [ 17]: Loss 0.03008
Validation:	Loss 0.02748	Accuracy 1.00000
Epoch [ 18]: Loss 0.02680
Validation:	Loss 0.02409	Accuracy 1.00000
Epoch [ 19]: Loss 0.02229
Validation:	Loss 0.01819	Accuracy 1.00000
Epoch [ 20]: Loss 0.01595
Validation:	Loss 0.01218	Accuracy 1.00000
Epoch [ 21]: Loss 0.01034
Validation:	Loss 0.00896	Accuracy 1.00000
Epoch [ 22]: Loss 0.00822
Validation:	Loss 0.00769	Accuracy 1.00000
Epoch [ 23]: Loss 0.00714
Validation:	Loss 0.00682	Accuracy 1.00000
Epoch [ 24]: Loss 0.00636
Validation:	Loss 0.00613	Accuracy 1.00000
Epoch [ 25]: Loss 0.00581
Validation:	Loss 0.00561	Accuracy 1.00000

We can also train the compact model with the exact same code!

ps_trained2, st_trained2 = main(SpiralClassifierCompact)

┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`.
└ @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:18
Epoch [  1]: Loss 0.77134
Validation:	Loss 0.60913	Accuracy 0.53125
Epoch [  2]: Loss 0.54649
Validation:	Loss 0.43598	Accuracy 1.00000
Epoch [  3]: Loss 0.37718
Validation:	Loss 0.30647	Accuracy 1.00000
Epoch [  4]: Loss 0.26789
Validation:	Loss 0.21552	Accuracy 1.00000
Epoch [  5]: Loss 0.18781
Validation:	Loss 0.15210	Accuracy 1.00000
Epoch [  6]: Loss 0.13776
Validation:	Loss 0.11700	Accuracy 1.00000
Epoch [  7]: Loss 0.10731
Validation:	Loss 0.09276	Accuracy 1.00000
Epoch [  8]: Loss 0.08600
Validation:	Loss 0.07551	Accuracy 1.00000
Epoch [  9]: Loss 0.07019
Validation:	Loss 0.06240	Accuracy 1.00000
Epoch [ 10]: Loss 0.05839
Validation:	Loss 0.05239	Accuracy 1.00000
Epoch [ 11]: Loss 0.04939
Validation:	Loss 0.04478	Accuracy 1.00000
Epoch [ 12]: Loss 0.04260
Validation:	Loss 0.03888	Accuracy 1.00000
Epoch [ 13]: Loss 0.03711
Validation:	Loss 0.03419	Accuracy 1.00000
Epoch [ 14]: Loss 0.03277
Validation:	Loss 0.03039	Accuracy 1.00000
Epoch [ 15]: Loss 0.02920
Validation:	Loss 0.02725	Accuracy 1.00000
Epoch [ 16]: Loss 0.02627
Validation:	Loss 0.02462	Accuracy 1.00000
Epoch [ 17]: Loss 0.02381
Validation:	Loss 0.02237	Accuracy 1.00000
Epoch [ 18]: Loss 0.02164
Validation:	Loss 0.02043	Accuracy 1.00000
Epoch [ 19]: Loss 0.01982
Validation:	Loss 0.01874	Accuracy 1.00000
Epoch [ 20]: Loss 0.01823
Validation:	Loss 0.01726	Accuracy 1.00000
Epoch [ 21]: Loss 0.01680
Validation:	Loss 0.01596	Accuracy 1.00000
Epoch [ 22]: Loss 0.01552
Validation:	Loss 0.01481	Accuracy 1.00000
Epoch [ 23]: Loss 0.01444
Validation:	Loss 0.01378	Accuracy 1.00000
Epoch [ 24]: Loss 0.01346
Validation:	Loss 0.01285	Accuracy 1.00000
Epoch [ 25]: Loss 0.01255
Validation:	Loss 0.01201	Accuracy 1.00000

Saving the Model

We can save the model using JLD2 (and any other serialization library of your choice) Note that we transfer the model to CPU before saving. Additionally, we recommend that you don't save the model struct and only save the parameters and states.

@save "trained_model.jld2" ps_trained st_trained

Let's try loading the model

@load "trained_model.jld2" ps_trained st_trained

2-element Vector{Symbol}:
 :ps_trained
 :st_trained

Appendix

using InteractiveUtils
InteractiveUtils.versioninfo()

if @isdefined(MLDataDevices)
    if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
        println()
        CUDA.versioninfo()
    end

    if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
        println()
        AMDGPU.versioninfo()
    end
end

Julia Version 1.12.5
Commit 5fe89b8ddc1 (2026-02-09 16:05 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, icelake-server)
  GC: Built with stock GC
Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores)
Environment:
  JULIA_DEBUG = Literate
  LD_LIBRARY_PATH = 
  JULIA_NUM_THREADS = 4
  JULIA_CPU_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0

---

This page was generated using Literate.jl.