Lux.jl Docs

Appearance


Trusted by

Scientific Computing

SciML.ai

Machine Learning

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 ​

julia
using ADTypes, Lux, LuxCUDA, JLD2, MLUtils, Optimisers, Zygote, Printf, Random, Statistics
Precompiling Lux...
    494.7 ms  ✓ Requires
    592.3 ms  ✓ CpuId
    505.2 ms  ✓ Compat
    590.8 ms  ✓ DocStringExtensions
    450.7 ms  ✓ JLLWrappers
    758.0 ms  ✓ Static
   2445.0 ms  ✓ MacroTools
    361.9 ms  ✓ Compat → CompatLinearAlgebraExt
   1023.6 ms  ✓ LazyArtifacts
    564.8 ms  ✓ LogExpFunctions
    582.4 ms  ✓ Hwloc_jll
    582.9 ms  ✓ OpenSpecFun_jll
    385.4 ms  ✓ BitTwiddlingConvenienceFunctions
    627.8 ms  ✓ CommonSubexpressions
    984.2 ms  ✓ CPUSummary
   1512.2 ms  ✓ DispatchDoctor
   1426.8 ms  ✓ Setfield
    569.8 ms  ✓ Functors
   1175.6 ms  ✓ ChainRulesCore
   1452.8 ms  ✓ StaticArrayInterface
   7339.0 ms  ✓ StaticArrays
   1365.9 ms  ✓ LLVMExtra_jll
    618.4 ms  ✓ PolyesterWeave
    407.5 ms  ✓ DispatchDoctor → DispatchDoctorEnzymeCoreExt
   2118.5 ms  ✓ Hwloc
   1171.6 ms  ✓ LuxCore
   2559.2 ms  ✓ SpecialFunctions
    804.5 ms  ✓ MLDataDevices
    388.6 ms  ✓ ADTypes → ADTypesChainRulesCoreExt
    403.0 ms  ✓ ArrayInterface → ArrayInterfaceChainRulesCoreExt
   1129.8 ms  ✓ Optimisers
    719.6 ms  ✓ DispatchDoctor → DispatchDoctorChainRulesCoreExt
    466.0 ms  ✓ CloseOpenIntervals
   1344.8 ms  ✓ LogExpFunctions → LogExpFunctionsChainRulesCoreExt
    569.5 ms  ✓ LayoutPointers
    623.9 ms  ✓ StaticArrays → StaticArraysChainRulesCoreExt
    648.5 ms  ✓ StaticArrays → StaticArraysStatisticsExt
    587.6 ms  ✓ ConstructionBase → ConstructionBaseStaticArraysExt
    567.9 ms  ✓ Adapt → AdaptStaticArraysExt
    648.9 ms  ✓ StaticArrayInterface → StaticArrayInterfaceStaticArraysExt
    575.2 ms  ✓ LuxCore → LuxCoreChainRulesCoreExt
    497.4 ms  ✓ LuxCore → LuxCoreFunctorsExt
    490.1 ms  ✓ LuxCore → LuxCoreEnzymeCoreExt
    460.6 ms  ✓ LuxCore → LuxCoreSetfieldExt
   1700.5 ms  ✓ SpecialFunctions → SpecialFunctionsChainRulesCoreExt
   2672.5 ms  ✓ WeightInitializers
    665.7 ms  ✓ MLDataDevices → MLDataDevicesChainRulesCoreExt
    467.8 ms  ✓ LuxCore → LuxCoreMLDataDevicesExt
    415.0 ms  ✓ Optimisers → OptimisersEnzymeCoreExt
    924.8 ms  ✓ StrideArraysCore
   6723.2 ms  ✓ LLVM
   3819.3 ms  ✓ ForwardDiff
    901.3 ms  ✓ WeightInitializers → WeightInitializersChainRulesCoreExt
    829.1 ms  ✓ Polyester
    884.8 ms  ✓ ForwardDiff → ForwardDiffStaticArraysExt
   1972.7 ms  ✓ UnsafeAtomicsLLVM
   4835.6 ms  ✓ KernelAbstractions
   1544.1 ms  ✓ KernelAbstractions → LinearAlgebraExt
   1605.3 ms  ✓ KernelAbstractions → EnzymeExt
   6109.4 ms  ✓ NNlib
   1740.8 ms  ✓ NNlib → NNlibEnzymeCoreExt
   1747.1 ms  ✓ NNlib → NNlibForwardDiffExt
   6504.7 ms  ✓ LuxLib
   9941.7 ms  ✓ Lux
  64 dependencies successfully precompiled in 53 seconds. 59 already precompiled.
Precompiling LuxCUDA...
    359.9 ms  ✓ Scratch
   1259.3 ms  ✓ SentinelArrays
    899.0 ms  ✓ CUDA_Driver_jll
   2662.4 ms  ✓ TimerOutputs
    573.5 ms  ✓ NVTX_jll
    560.5 ms  ✓ demumble_jll
    539.1 ms  ✓ JuliaNVTXCallbacks_jll
   3676.6 ms  ✓ Test
   1688.9 ms  ✓ DataStructures
   1974.1 ms  ✓ StringManipulation
   2190.2 ms  ✓ GPUArrays
   1857.7 ms  ✓ KernelAbstractions → SparseArraysExt
   2763.2 ms  ✓ CUDA_Runtime_jll
   1309.9 ms  ✓ NVTX
    499.5 ms  ✓ BFloat16s
    483.2 ms  ✓ SortingAlgorithms
   1314.4 ms  ✓ AbstractFFTs → AbstractFFTsTestExt
   2239.2 ms  ✓ CUDNN_jll
   1383.0 ms  ✓ LLVM → BFloat16sExt
  20121.5 ms  ✓ PrettyTables
  27228.2 ms  ✓ GPUCompiler
  45372.2 ms  ✓ DataFrames
  52700.1 ms  ✓ CUDA
   5660.8 ms  ✓ Atomix → AtomixCUDAExt
   8816.0 ms  ✓ cuDNN
   5889.2 ms  ✓ LuxCUDA
  26 dependencies successfully precompiled in 148 seconds. 74 already precompiled.
Precompiling MLDataDevicesGPUArraysExt...
   1532.8 ms  ✓ MLDataDevices → MLDataDevicesGPUArraysExt
  1 dependency successfully precompiled in 2 seconds. 42 already precompiled.
Precompiling WeightInitializersGPUArraysExt...
   1509.4 ms  ✓ WeightInitializers → WeightInitializersGPUArraysExt
  1 dependency successfully precompiled in 2 seconds. 46 already precompiled.
Precompiling ChainRulesCoreSparseArraysExt...
    631.9 ms  ✓ ChainRulesCore → ChainRulesCoreSparseArraysExt
  1 dependency successfully precompiled in 1 seconds. 11 already precompiled.
Precompiling MLDataDevicesSparseArraysExt...
    668.5 ms  ✓ MLDataDevices → MLDataDevicesSparseArraysExt
  1 dependency successfully precompiled in 1 seconds. 17 already precompiled.
Precompiling AbstractFFTsChainRulesCoreExt...
    416.4 ms  ✓ AbstractFFTs → AbstractFFTsChainRulesCoreExt
  1 dependency successfully precompiled in 0 seconds. 9 already precompiled.
Precompiling ArrayInterfaceCUDAExt...
   5096.8 ms  ✓ ArrayInterface → ArrayInterfaceCUDAExt
  1 dependency successfully precompiled in 5 seconds. 101 already precompiled.
Precompiling NNlibCUDAExt...
   5305.3 ms  ✓ CUDA → ChainRulesCoreExt
   5643.7 ms  ✓ NNlib → NNlibCUDAExt
  2 dependencies successfully precompiled in 6 seconds. 102 already precompiled.
Precompiling MLDataDevicesCUDAExt...
   5099.6 ms  ✓ MLDataDevices → MLDataDevicesCUDAExt
  1 dependency successfully precompiled in 5 seconds. 104 already precompiled.
Precompiling LuxLibCUDAExt...
   5597.0 ms  ✓ CUDA → SpecialFunctionsExt
   5646.9 ms  ✓ CUDA → EnzymeCoreExt
   6313.4 ms  ✓ LuxLib → LuxLibCUDAExt
  3 dependencies successfully precompiled in 7 seconds. 167 already precompiled.
Precompiling WeightInitializersCUDAExt...
   5658.6 ms  ✓ WeightInitializers → WeightInitializersCUDAExt
  1 dependency successfully precompiled in 6 seconds. 109 already precompiled.
Precompiling NNlibCUDACUDNNExt...
   5694.9 ms  ✓ NNlib → NNlibCUDACUDNNExt
  1 dependency successfully precompiled in 6 seconds. 106 already precompiled.
Precompiling MLDataDevicescuDNNExt...
   5434.0 ms  ✓ MLDataDevices → MLDataDevicescuDNNExt
  1 dependency successfully precompiled in 6 seconds. 107 already precompiled.
Precompiling LuxLibcuDNNExt...
   6345.7 ms  ✓ LuxLib → LuxLibcuDNNExt
  1 dependency successfully precompiled in 7 seconds. 174 already precompiled.
Precompiling JLD2...
   4466.5 ms  ✓ FileIO
  33349.5 ms  ✓ JLD2
  2 dependencies successfully precompiled in 38 seconds. 30 already precompiled.
Precompiling MLUtils...
    390.2 ms  ✓ InverseFunctions → InverseFunctionsDatesExt
    425.5 ms  ✓ ContextVariablesX
   1161.5 ms  ✓ SimpleTraits
    589.7 ms  ✓ InverseFunctions → InverseFunctionsTestExt
   1125.0 ms  ✓ SplittablesBase
    428.9 ms  ✓ LogExpFunctions → LogExpFunctionsInverseFunctionsExt
    588.0 ms  ✓ FLoopsBase
   2984.0 ms  ✓ Accessors
    705.6 ms  ✓ Accessors → AccessorsTestExt
    831.2 ms  ✓ Accessors → AccessorsDatesExt
   2316.0 ms  ✓ StatsBase
    795.3 ms  ✓ BangBang
    706.4 ms  ✓ Accessors → AccessorsStaticArraysExt
    501.7 ms  ✓ BangBang → BangBangChainRulesCoreExt
    520.5 ms  ✓ BangBang → BangBangTablesExt
    720.2 ms  ✓ BangBang → BangBangStaticArraysExt
    841.3 ms  ✓ MicroCollections
   2704.9 ms  ✓ Transducers
    630.5 ms  ✓ Transducers → TransducersAdaptExt
   5113.2 ms  ✓ FLoops
   7489.8 ms  ✓ MLUtils
  21 dependencies successfully precompiled in 21 seconds. 91 already precompiled.
Precompiling BangBangDataFramesExt...
   1621.7 ms  ✓ BangBang → BangBangDataFramesExt
  1 dependency successfully precompiled in 2 seconds. 46 already precompiled.
Precompiling TransducersDataFramesExt...
   1378.5 ms  ✓ Transducers → TransducersDataFramesExt
  1 dependency successfully precompiled in 2 seconds. 61 already precompiled.
Precompiling MLDataDevicesMLUtilsExt...
   2525.6 ms  ✓ MLDataDevices → MLDataDevicesMLUtilsExt
  1 dependency successfully precompiled in 3 seconds. 116 already precompiled.
Precompiling LuxMLUtilsExt...
   3234.6 ms  ✓ Lux → LuxMLUtilsExt
  1 dependency successfully precompiled in 4 seconds. 178 already precompiled.
Precompiling Zygote...
    948.7 ms  ✓ ZygoteRules
   1871.1 ms  ✓ IRTools
   5590.1 ms  ✓ ChainRules
  34780.0 ms  ✓ Zygote
  4 dependencies successfully precompiled in 41 seconds. 82 already precompiled.
Precompiling AccessorsStructArraysExt...
    507.2 ms  ✓ Accessors → AccessorsStructArraysExt
  1 dependency successfully precompiled in 1 seconds. 16 already precompiled.
Precompiling BangBangStructArraysExt...
    484.3 ms  ✓ BangBang → BangBangStructArraysExt
  1 dependency successfully precompiled in 1 seconds. 22 already precompiled.
Precompiling StructArraysStaticArraysExt...
    654.9 ms  ✓ StructArrays → StructArraysStaticArraysExt
  1 dependency successfully precompiled in 1 seconds. 18 already precompiled.
Precompiling ArrayInterfaceChainRulesExt...
    770.0 ms  ✓ ArrayInterface → ArrayInterfaceChainRulesExt
  1 dependency successfully precompiled in 1 seconds. 39 already precompiled.
Precompiling MLDataDevicesChainRulesExt...
    819.8 ms  ✓ MLDataDevices → MLDataDevicesChainRulesExt
  1 dependency successfully precompiled in 1 seconds. 40 already precompiled.
Precompiling MLDataDevicesFillArraysExt...
    447.6 ms  ✓ MLDataDevices → MLDataDevicesFillArraysExt
  1 dependency successfully precompiled in 0 seconds. 15 already precompiled.
Precompiling MLDataDevicesZygoteExt...
   1650.1 ms  ✓ MLDataDevices → MLDataDevicesZygoteExt
  1 dependency successfully precompiled in 2 seconds. 93 already precompiled.
Precompiling LuxZygoteExt...
   3497.9 ms  ✓ Lux → LuxZygoteExt
  1 dependency successfully precompiled in 4 seconds. 162 already precompiled.
Precompiling ZygoteColorsExt...
   1776.1 ms  ✓ Zygote → ZygoteColorsExt
  1 dependency successfully precompiled in 2 seconds. 89 already precompiled.

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.

julia
function get_dataloaders(; 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))
    # 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),
        # Don't shuffle the validation data
        DataLoader(collect.((x_val, y_val)); batchsize=128, shuffle=false))
end
get_dataloaders (generic function with 1 method)

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 fieldnames 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.

julia
struct SpiralClassifier{L, C} <: Lux.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.

julia
function SpiralClassifier(in_dims, hidden_dims, out_dims)
    return SpiralClassifier(
        LSTMCell(in_dims => hidden_dims), Dense(hidden_dims => out_dims, sigmoid))
end
Main.var"##230".SpiralClassifier

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.

julia
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

julia
function SpiralClassifierCompact(in_dims, hidden_dims, out_dims)
    lstm_cell = LSTMCell(in_dims => hidden_dims)
    classifier = Dense(hidden_dims => out_dims, sigmoid)
    return @compact(; lstm_cell, classifier) 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
SpiralClassifierCompact (generic function with 1 method)

Defining Accuracy, Loss and Optimiser ​

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

julia
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)
accuracy (generic function with 1 method)

Training the Model ​

julia
function main(model_type)
    dev = gpu_device()

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

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

    train_state = Training.TrainState(model, ps, st, Adam(0.01f0))

    for epoch in 1:25
        # Train the model
        for (x, y) in train_loader
            (_, loss, _, train_state) = Training.single_train_step!(
                AutoZygote(), lossfn, (x, y), train_state)

            @printf "Epoch [%3d]: Loss %4.5f\n" epoch loss
        end

        # Validate the model
        st_ = Lux.testmode(train_state.states)
        for (x, y) in val_loader
            ŷ, st_ = model(x, train_state.parameters, st_)
            loss = lossfn(ŷ, y)
            acc = accuracy(ŷ, y)
            @printf "Validation: Loss %4.5f Accuracy %4.5f\n" loss acc
        end
    end

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

ps_trained, st_trained = main(SpiralClassifier)
Epoch [  1]: Loss 0.61367
Epoch [  1]: Loss 0.59712
Epoch [  1]: Loss 0.56119
Epoch [  1]: Loss 0.54298
Epoch [  1]: Loss 0.51899
Epoch [  1]: Loss 0.50466
Epoch [  1]: Loss 0.47137
Validation: Loss 0.46711 Accuracy 1.00000
Validation: Loss 0.47555 Accuracy 1.00000
Epoch [  2]: Loss 0.46574
Epoch [  2]: Loss 0.45492
Epoch [  2]: Loss 0.43776
Epoch [  2]: Loss 0.42556
Epoch [  2]: Loss 0.41798
Epoch [  2]: Loss 0.38898
Epoch [  2]: Loss 0.39934
Validation: Loss 0.36997 Accuracy 1.00000
Validation: Loss 0.37992 Accuracy 1.00000
Epoch [  3]: Loss 0.37373
Epoch [  3]: Loss 0.36598
Epoch [  3]: Loss 0.33692
Epoch [  3]: Loss 0.32881
Epoch [  3]: Loss 0.30676
Epoch [  3]: Loss 0.31733
Epoch [  3]: Loss 0.32936
Validation: Loss 0.28536 Accuracy 1.00000
Validation: Loss 0.29610 Accuracy 1.00000
Epoch [  4]: Loss 0.28156
Epoch [  4]: Loss 0.27879
Epoch [  4]: Loss 0.25902
Epoch [  4]: Loss 0.25110
Epoch [  4]: Loss 0.25401
Epoch [  4]: Loss 0.23588
Epoch [  4]: Loss 0.23325
Validation: Loss 0.21578 Accuracy 1.00000
Validation: Loss 0.22649 Accuracy 1.00000
Epoch [  5]: Loss 0.22262
Epoch [  5]: Loss 0.21548
Epoch [  5]: Loss 0.18405
Epoch [  5]: Loss 0.18794
Epoch [  5]: Loss 0.19170
Epoch [  5]: Loss 0.17481
Epoch [  5]: Loss 0.15956
Validation: Loss 0.16018 Accuracy 1.00000
Validation: Loss 0.17008 Accuracy 1.00000
Epoch [  6]: Loss 0.16037
Epoch [  6]: Loss 0.14633
Epoch [  6]: Loss 0.13868
Epoch [  6]: Loss 0.15006
Epoch [  6]: Loss 0.13592
Epoch [  6]: Loss 0.13859
Epoch [  6]: Loss 0.11643
Validation: Loss 0.11760 Accuracy 1.00000
Validation: Loss 0.12588 Accuracy 1.00000
Epoch [  7]: Loss 0.11172
Epoch [  7]: Loss 0.11035
Epoch [  7]: Loss 0.10630
Epoch [  7]: Loss 0.11381
Epoch [  7]: Loss 0.10564
Epoch [  7]: Loss 0.09000
Epoch [  7]: Loss 0.07393
Validation: Loss 0.08442 Accuracy 1.00000
Validation: Loss 0.09060 Accuracy 1.00000
Epoch [  8]: Loss 0.08092
Epoch [  8]: Loss 0.08684
Epoch [  8]: Loss 0.07771
Epoch [  8]: Loss 0.07386
Epoch [  8]: Loss 0.06715
Epoch [  8]: Loss 0.06483
Epoch [  8]: Loss 0.06035
Validation: Loss 0.05895 Accuracy 1.00000
Validation: Loss 0.06310 Accuracy 1.00000
Epoch [  9]: Loss 0.05952
Epoch [  9]: Loss 0.05917
Epoch [  9]: Loss 0.05295
Epoch [  9]: Loss 0.05165
Epoch [  9]: Loss 0.05267
Epoch [  9]: Loss 0.04426
Epoch [  9]: Loss 0.03650
Validation: Loss 0.04361 Accuracy 1.00000
Validation: Loss 0.04654 Accuracy 1.00000
Epoch [ 10]: Loss 0.04547
Epoch [ 10]: Loss 0.04390
Epoch [ 10]: Loss 0.04146
Epoch [ 10]: Loss 0.04115
Epoch [ 10]: Loss 0.03283
Epoch [ 10]: Loss 0.03817
Epoch [ 10]: Loss 0.03660
Validation: Loss 0.03522 Accuracy 1.00000
Validation: Loss 0.03762 Accuracy 1.00000
Epoch [ 11]: Loss 0.03587
Epoch [ 11]: Loss 0.03404
Epoch [ 11]: Loss 0.03433
Epoch [ 11]: Loss 0.03436
Epoch [ 11]: Loss 0.03085
Epoch [ 11]: Loss 0.03042
Epoch [ 11]: Loss 0.03078
Validation: Loss 0.02986 Accuracy 1.00000
Validation: Loss 0.03195 Accuracy 1.00000
Epoch [ 12]: Loss 0.03198
Epoch [ 12]: Loss 0.02974
Epoch [ 12]: Loss 0.02867
Epoch [ 12]: Loss 0.02781
Epoch [ 12]: Loss 0.02750
Epoch [ 12]: Loss 0.02554
Epoch [ 12]: Loss 0.02669
Validation: Loss 0.02595 Accuracy 1.00000
Validation: Loss 0.02781 Accuracy 1.00000
Epoch [ 13]: Loss 0.02551
Epoch [ 13]: Loss 0.02479
Epoch [ 13]: Loss 0.02459
Epoch [ 13]: Loss 0.02633
Epoch [ 13]: Loss 0.02437
Epoch [ 13]: Loss 0.02461
Epoch [ 13]: Loss 0.02105
Validation: Loss 0.02294 Accuracy 1.00000
Validation: Loss 0.02461 Accuracy 1.00000
Epoch [ 14]: Loss 0.02234
Epoch [ 14]: Loss 0.02175
Epoch [ 14]: Loss 0.02281
Epoch [ 14]: Loss 0.02226
Epoch [ 14]: Loss 0.02327
Epoch [ 14]: Loss 0.02040
Epoch [ 14]: Loss 0.02104
Validation: Loss 0.02052 Accuracy 1.00000
Validation: Loss 0.02205 Accuracy 1.00000
Epoch [ 15]: Loss 0.01976
Epoch [ 15]: Loss 0.01968
Epoch [ 15]: Loss 0.02029
Epoch [ 15]: Loss 0.01943
Epoch [ 15]: Loss 0.02006
Epoch [ 15]: Loss 0.01961
Epoch [ 15]: Loss 0.02037
Validation: Loss 0.01851 Accuracy 1.00000
Validation: Loss 0.01993 Accuracy 1.00000
Epoch [ 16]: Loss 0.01660
Epoch [ 16]: Loss 0.01825
Epoch [ 16]: Loss 0.01938
Epoch [ 16]: Loss 0.01776
Epoch [ 16]: Loss 0.01784
Epoch [ 16]: Loss 0.01812
Epoch [ 16]: Loss 0.01656
Validation: Loss 0.01681 Accuracy 1.00000
Validation: Loss 0.01812 Accuracy 1.00000
Epoch [ 17]: Loss 0.01682
Epoch [ 17]: Loss 0.01837
Epoch [ 17]: Loss 0.01527
Epoch [ 17]: Loss 0.01523
Epoch [ 17]: Loss 0.01565
Epoch [ 17]: Loss 0.01684
Epoch [ 17]: Loss 0.01571
Validation: Loss 0.01537 Accuracy 1.00000
Validation: Loss 0.01659 Accuracy 1.00000
Epoch [ 18]: Loss 0.01446
Epoch [ 18]: Loss 0.01608
Epoch [ 18]: Loss 0.01456
Epoch [ 18]: Loss 0.01542
Epoch [ 18]: Loss 0.01393
Epoch [ 18]: Loss 0.01548
Epoch [ 18]: Loss 0.01444
Validation: Loss 0.01414 Accuracy 1.00000
Validation: Loss 0.01527 Accuracy 1.00000
Epoch [ 19]: Loss 0.01414
Epoch [ 19]: Loss 0.01495
Epoch [ 19]: Loss 0.01299
Epoch [ 19]: Loss 0.01421
Epoch [ 19]: Loss 0.01391
Epoch [ 19]: Loss 0.01216
Epoch [ 19]: Loss 0.01544
Validation: Loss 0.01306 Accuracy 1.00000
Validation: Loss 0.01412 Accuracy 1.00000
Epoch [ 20]: Loss 0.01372
Epoch [ 20]: Loss 0.01304
Epoch [ 20]: Loss 0.01355
Epoch [ 20]: Loss 0.01263
Epoch [ 20]: Loss 0.01212
Epoch [ 20]: Loss 0.01175
Epoch [ 20]: Loss 0.01186
Validation: Loss 0.01209 Accuracy 1.00000
Validation: Loss 0.01307 Accuracy 1.00000
Epoch [ 21]: Loss 0.01197
Epoch [ 21]: Loss 0.01246
Epoch [ 21]: Loss 0.01219
Epoch [ 21]: Loss 0.01205
Epoch [ 21]: Loss 0.01180
Epoch [ 21]: Loss 0.01064
Epoch [ 21]: Loss 0.01099
Validation: Loss 0.01119 Accuracy 1.00000
Validation: Loss 0.01209 Accuracy 1.00000
Epoch [ 22]: Loss 0.01062
Epoch [ 22]: Loss 0.01161
Epoch [ 22]: Loss 0.01181
Epoch [ 22]: Loss 0.01134
Epoch [ 22]: Loss 0.00982
Epoch [ 22]: Loss 0.01042
Epoch [ 22]: Loss 0.01010
Validation: Loss 0.01026 Accuracy 1.00000
Validation: Loss 0.01107 Accuracy 1.00000
Epoch [ 23]: Loss 0.01015
Epoch [ 23]: Loss 0.00991
Epoch [ 23]: Loss 0.00929
Epoch [ 23]: Loss 0.00944
Epoch [ 23]: Loss 0.01016
Epoch [ 23]: Loss 0.01036
Epoch [ 23]: Loss 0.01054
Validation: Loss 0.00921 Accuracy 1.00000
Validation: Loss 0.00992 Accuracy 1.00000
Epoch [ 24]: Loss 0.00984
Epoch [ 24]: Loss 0.00895
Epoch [ 24]: Loss 0.00880
Epoch [ 24]: Loss 0.00834
Epoch [ 24]: Loss 0.00883
Epoch [ 24]: Loss 0.00829
Epoch [ 24]: Loss 0.00897
Validation: Loss 0.00816 Accuracy 1.00000
Validation: Loss 0.00877 Accuracy 1.00000
Epoch [ 25]: Loss 0.00785
Epoch [ 25]: Loss 0.00789
Epoch [ 25]: Loss 0.00761
Epoch [ 25]: Loss 0.00844
Epoch [ 25]: Loss 0.00776
Epoch [ 25]: Loss 0.00756
Epoch [ 25]: Loss 0.00837
Validation: Loss 0.00738 Accuracy 1.00000
Validation: Loss 0.00791 Accuracy 1.00000

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

julia
ps_trained2, st_trained2 = main(SpiralClassifierCompact)
Epoch [  1]: Loss 0.63487
Epoch [  1]: Loss 0.60560
Epoch [  1]: Loss 0.56825
Epoch [  1]: Loss 0.53385
Epoch [  1]: Loss 0.51626
Epoch [  1]: Loss 0.50272
Epoch [  1]: Loss 0.48309
Validation: Loss 0.46197 Accuracy 1.00000
Validation: Loss 0.45618 Accuracy 1.00000
Epoch [  2]: Loss 0.47915
Epoch [  2]: Loss 0.46448
Epoch [  2]: Loss 0.45141
Epoch [  2]: Loss 0.42117
Epoch [  2]: Loss 0.41135
Epoch [  2]: Loss 0.38516
Epoch [  2]: Loss 0.37026
Validation: Loss 0.36330 Accuracy 1.00000
Validation: Loss 0.35627 Accuracy 1.00000
Epoch [  3]: Loss 0.37922
Epoch [  3]: Loss 0.35917
Epoch [  3]: Loss 0.34686
Epoch [  3]: Loss 0.32908
Epoch [  3]: Loss 0.32395
Epoch [  3]: Loss 0.31569
Epoch [  3]: Loss 0.28576
Validation: Loss 0.27709 Accuracy 1.00000
Validation: Loss 0.26881 Accuracy 1.00000
Epoch [  4]: Loss 0.30241
Epoch [  4]: Loss 0.26503
Epoch [  4]: Loss 0.26196
Epoch [  4]: Loss 0.26078
Epoch [  4]: Loss 0.24990
Epoch [  4]: Loss 0.23581
Epoch [  4]: Loss 0.21976
Validation: Loss 0.20740 Accuracy 1.00000
Validation: Loss 0.19881 Accuracy 1.00000
Epoch [  5]: Loss 0.22857
Epoch [  5]: Loss 0.21933
Epoch [  5]: Loss 0.20009
Epoch [  5]: Loss 0.18151
Epoch [  5]: Loss 0.17477
Epoch [  5]: Loss 0.18024
Epoch [  5]: Loss 0.18392
Validation: Loss 0.15292 Accuracy 1.00000
Validation: Loss 0.14489 Accuracy 1.00000
Epoch [  6]: Loss 0.15658
Epoch [  6]: Loss 0.15041
Epoch [  6]: Loss 0.15387
Epoch [  6]: Loss 0.14809
Epoch [  6]: Loss 0.13858
Epoch [  6]: Loss 0.13707
Epoch [  6]: Loss 0.11415
Validation: Loss 0.11155 Accuracy 1.00000
Validation: Loss 0.10480 Accuracy 1.00000
Epoch [  7]: Loss 0.11169
Epoch [  7]: Loss 0.11783
Epoch [  7]: Loss 0.09649
Epoch [  7]: Loss 0.11305
Epoch [  7]: Loss 0.10724
Epoch [  7]: Loss 0.09856
Epoch [  7]: Loss 0.08203
Validation: Loss 0.07984 Accuracy 1.00000
Validation: Loss 0.07474 Accuracy 1.00000
Epoch [  8]: Loss 0.08615
Epoch [  8]: Loss 0.08119
Epoch [  8]: Loss 0.07777
Epoch [  8]: Loss 0.07642
Epoch [  8]: Loss 0.06875
Epoch [  8]: Loss 0.06414
Epoch [  8]: Loss 0.07095
Validation: Loss 0.05577 Accuracy 1.00000
Validation: Loss 0.05238 Accuracy 1.00000
Epoch [  9]: Loss 0.06346
Epoch [  9]: Loss 0.05796
Epoch [  9]: Loss 0.05713
Epoch [  9]: Loss 0.05191
Epoch [  9]: Loss 0.04658
Epoch [  9]: Loss 0.04425
Epoch [  9]: Loss 0.04977
Validation: Loss 0.04153 Accuracy 1.00000
Validation: Loss 0.03919 Accuracy 1.00000
Epoch [ 10]: Loss 0.04668
Epoch [ 10]: Loss 0.04291
Epoch [ 10]: Loss 0.04354
Epoch [ 10]: Loss 0.03695
Epoch [ 10]: Loss 0.03875
Epoch [ 10]: Loss 0.03698
Epoch [ 10]: Loss 0.03998
Validation: Loss 0.03359 Accuracy 1.00000
Validation: Loss 0.03171 Accuracy 1.00000
Epoch [ 11]: Loss 0.03982
Epoch [ 11]: Loss 0.03447
Epoch [ 11]: Loss 0.03164
Epoch [ 11]: Loss 0.03227
Epoch [ 11]: Loss 0.03393
Epoch [ 11]: Loss 0.03081
Epoch [ 11]: Loss 0.03278
Validation: Loss 0.02848 Accuracy 1.00000
Validation: Loss 0.02685 Accuracy 1.00000
Epoch [ 12]: Loss 0.03180
Epoch [ 12]: Loss 0.03115
Epoch [ 12]: Loss 0.02951
Epoch [ 12]: Loss 0.02772
Epoch [ 12]: Loss 0.02839
Epoch [ 12]: Loss 0.02640
Epoch [ 12]: Loss 0.02476
Validation: Loss 0.02476 Accuracy 1.00000
Validation: Loss 0.02332 Accuracy 1.00000
Epoch [ 13]: Loss 0.02545
Epoch [ 13]: Loss 0.02697
Epoch [ 13]: Loss 0.02738
Epoch [ 13]: Loss 0.02455
Epoch [ 13]: Loss 0.02388
Epoch [ 13]: Loss 0.02414
Epoch [ 13]: Loss 0.02516
Validation: Loss 0.02190 Accuracy 1.00000
Validation: Loss 0.02060 Accuracy 1.00000
Epoch [ 14]: Loss 0.02352
Epoch [ 14]: Loss 0.02257
Epoch [ 14]: Loss 0.02298
Epoch [ 14]: Loss 0.02145
Epoch [ 14]: Loss 0.02305
Epoch [ 14]: Loss 0.02212
Epoch [ 14]: Loss 0.02179
Validation: Loss 0.01961 Accuracy 1.00000
Validation: Loss 0.01842 Accuracy 1.00000
Epoch [ 15]: Loss 0.02071
Epoch [ 15]: Loss 0.01987
Epoch [ 15]: Loss 0.02049
Epoch [ 15]: Loss 0.02108
Epoch [ 15]: Loss 0.01990
Epoch [ 15]: Loss 0.02014
Epoch [ 15]: Loss 0.01900
Validation: Loss 0.01770 Accuracy 1.00000
Validation: Loss 0.01660 Accuracy 1.00000
Epoch [ 16]: Loss 0.01941
Epoch [ 16]: Loss 0.01881
Epoch [ 16]: Loss 0.01938
Epoch [ 16]: Loss 0.01864
Epoch [ 16]: Loss 0.01674
Epoch [ 16]: Loss 0.01716
Epoch [ 16]: Loss 0.01982
Validation: Loss 0.01608 Accuracy 1.00000
Validation: Loss 0.01506 Accuracy 1.00000
Epoch [ 17]: Loss 0.01845
Epoch [ 17]: Loss 0.01507
Epoch [ 17]: Loss 0.01690
Epoch [ 17]: Loss 0.01715
Epoch [ 17]: Loss 0.01617
Epoch [ 17]: Loss 0.01686
Epoch [ 17]: Loss 0.01696
Validation: Loss 0.01468 Accuracy 1.00000
Validation: Loss 0.01373 Accuracy 1.00000
Epoch [ 18]: Loss 0.01689
Epoch [ 18]: Loss 0.01580
Epoch [ 18]: Loss 0.01489
Epoch [ 18]: Loss 0.01557
Epoch [ 18]: Loss 0.01535
Epoch [ 18]: Loss 0.01435
Epoch [ 18]: Loss 0.01281
Validation: Loss 0.01345 Accuracy 1.00000
Validation: Loss 0.01256 Accuracy 1.00000
Epoch [ 19]: Loss 0.01509
Epoch [ 19]: Loss 0.01547
Epoch [ 19]: Loss 0.01367
Epoch [ 19]: Loss 0.01350
Epoch [ 19]: Loss 0.01386
Epoch [ 19]: Loss 0.01397
Epoch [ 19]: Loss 0.01073
Validation: Loss 0.01237 Accuracy 1.00000
Validation: Loss 0.01153 Accuracy 1.00000
Epoch [ 20]: Loss 0.01276
Epoch [ 20]: Loss 0.01374
Epoch [ 20]: Loss 0.01321
Epoch [ 20]: Loss 0.01267
Epoch [ 20]: Loss 0.01399
Epoch [ 20]: Loss 0.01212
Epoch [ 20]: Loss 0.01078
Validation: Loss 0.01137 Accuracy 1.00000
Validation: Loss 0.01059 Accuracy 1.00000
Epoch [ 21]: Loss 0.01278
Epoch [ 21]: Loss 0.01214
Epoch [ 21]: Loss 0.01210
Epoch [ 21]: Loss 0.01149
Epoch [ 21]: Loss 0.01101
Epoch [ 21]: Loss 0.01205
Epoch [ 21]: Loss 0.01113
Validation: Loss 0.01033 Accuracy 1.00000
Validation: Loss 0.00961 Accuracy 1.00000
Epoch [ 22]: Loss 0.01214
Epoch [ 22]: Loss 0.01152
Epoch [ 22]: Loss 0.01091
Epoch [ 22]: Loss 0.01053
Epoch [ 22]: Loss 0.01014
Epoch [ 22]: Loss 0.00948
Epoch [ 22]: Loss 0.00916
Validation: Loss 0.00919 Accuracy 1.00000
Validation: Loss 0.00856 Accuracy 1.00000
Epoch [ 23]: Loss 0.01018
Epoch [ 23]: Loss 0.00982
Epoch [ 23]: Loss 0.01000
Epoch [ 23]: Loss 0.00905
Epoch [ 23]: Loss 0.00923
Epoch [ 23]: Loss 0.00874
Epoch [ 23]: Loss 0.00878
Validation: Loss 0.00817 Accuracy 1.00000
Validation: Loss 0.00764 Accuracy 1.00000
Epoch [ 24]: Loss 0.00863
Epoch [ 24]: Loss 0.00819
Epoch [ 24]: Loss 0.00890
Epoch [ 24]: Loss 0.00870
Epoch [ 24]: Loss 0.00836
Epoch [ 24]: Loss 0.00812
Epoch [ 24]: Loss 0.00795
Validation: Loss 0.00743 Accuracy 1.00000
Validation: Loss 0.00697 Accuracy 1.00000
Epoch [ 25]: Loss 0.00734
Epoch [ 25]: Loss 0.00837
Epoch [ 25]: Loss 0.00800
Epoch [ 25]: Loss 0.00784
Epoch [ 25]: Loss 0.00736
Epoch [ 25]: Loss 0.00760
Epoch [ 25]: Loss 0.00761
Validation: Loss 0.00688 Accuracy 1.00000
Validation: Loss 0.00647 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.

julia
@save "trained_model.jld2" ps_trained st_trained

Let's try loading the model

julia
@load "trained_model.jld2" ps_trained st_trained
2-element Vector{Symbol}:
 :ps_trained
 :st_trained

Appendix ​

julia
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.11.2
Commit 5e9a32e7af2 (2024-12-01 20:02 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 48 × AMD EPYC 7402 24-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver2)
Threads: 48 default, 0 interactive, 24 GC (on 2 virtual cores)
Environment:
  JULIA_CPU_THREADS = 2
  JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6
  LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
  JULIA_PKG_SERVER = 
  JULIA_NUM_THREADS = 48
  JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0
  JULIA_DEBUG = Literate

CUDA runtime 12.6, artifact installation
CUDA driver 12.6
NVIDIA driver 560.35.3

CUDA libraries: 
- CUBLAS: 12.6.4
- CURAND: 10.3.7
- CUFFT: 11.3.0
- CUSOLVER: 11.7.1
- CUSPARSE: 12.5.4
- CUPTI: 2024.3.2 (API 24.0.0)
- NVML: 12.0.0+560.35.3

Julia packages: 
- CUDA: 5.5.2
- CUDA_Driver_jll: 0.10.4+0
- CUDA_Runtime_jll: 0.15.5+0

Toolchain:
- Julia: 1.11.2
- LLVM: 16.0.6

Environment:
- JULIA_CUDA_HARD_MEMORY_LIMIT: 100%

2 devices:
  0: Quadro RTX 5000 (sm_75, 14.883 GiB / 16.000 GiB available)
  1: Quadro RTX 5000 (sm_75, 15.549 GiB / 16.000 GiB available)

This page was generated using Literate.jl.