Flux.jl

Flux is a pure-Julia machine learning library that provides lightweight abstractions for GPU and Automatic Differentiation (AD) support. You can define models using standard Julia closures (parameterized functions) or by composing built-in layers using Chain.

To train a model, use Flux.setup to initialize the optimizer state and Flux.train! to execute the training loop.

using Flux
data = [(x, 2x-x^3) for x in -2:0.1f0:2]

model = let
  w, b, v = (randn(Float32, 23) for _ in 1:3)  # parameters
  x -> sum(v .* tanh.(w*x .+ b))               # callable
end
# model = Chain(vcat, Dense(1 => 23, tanh), Dense(23 => 1, bias=false), only)

opt_state = Flux.setup(Adam(), model)
for epoch in 1:100
  Flux.train!((m,x,y) -> (m(x) - y)^2, model, data, opt_state)
end

using Plots
plot(x -> 2x-x^3, -2, 2, label="truth")
scatter!(model, -2:0.1f0:2, label="learned")

When training, images and labels should be grouped into mini-batches. For image data, this typically involves creating a 4D array of shape (height, width, channels, batch_size). Labels should be converted to a one-hot encoded format using Flux.onehotbatch.

function make_minibatch(X, Y, idxs)
   X_batch = Array{Float32}(undef, size(X)[1:end-1]..., 1, length(idxs))
   for i in 1:length(idxs)
       X_batch[:, :, :, i] = Float32.(X[:,:,idxs[i]])
   end
   Y_batch = onehotbatch(Y[idxs], 0:9)
   return (X_batch, Y_batch)
end

Optimisation rules (like Descent, Momentum, or Adam) define how parameters are adjusted.

1. Setup: Always call Flux.setup(rule, model) before training. This creates the necessary state (e.g., momentum buffers) for the specific model architecture.
2. Update: Use Flux.update!(opt_state, model, grads[1]) to apply the rule to the model parameters using the computed gradients.

# Initialise momentum
opt_state = Flux.setup(Momentum(0.01, 0.9), model)

for data in train_set
  grads = Flux.gradient(model) do m
      loss(m(data[1]), data[2])
  end

  # Update both model parameters and optimiser state
  Flux.update!(opt_state, model, grads[1])
end

You can manage specific GPU devices using gpu_device(id). Note that while CUDA.devices() uses 0-based indexing, gpu_device expects 1-based indexing. Once a device handle is obtained, you can move models or arrays to that device using the |> operator.

Note: Data movement across different devices is currently only supported for CUDA and AMDGPU backends; it is not supported for Metal.jl.

using Flux, CUDA;

# List devices (0-indexed)
CUDA.devices()

# Select device 0 using 1-based indexing
device0 = gpu_device(1)

# Move a model to the device
dense_model = Dense(2 => 3)
dense_model = dense_model |> device0

# Verify the device
CUDA.device(dense_model.weight)

You can persist and restore model weights using Flux.params in conjunction with the BSON package.

To save: Use BSON.@save to store the parameters extracted via params(model). It is recommended to move parameters to the CPU using cpu.(params(model)) before saving if you are training on a GPU.

To load:

1. Reconstruct the model architecture with the same dimensions used during training.
2. Load the saved parameters using BSON.@load.
3. Apply the loaded parameters to the model instance using Flux.loadparams!(model, loaded_params).

Flux models are parameterized functions. While parameters can be global variables or explicit arguments, Flux's preferred pattern is to encapsulate parameters within the function itself using closures or struct instances. This encapsulation allows for easy function composition using the Julia composition operator ∘.

Any callable struct that stores parameters can act as a valid Flux model.

struct Poly3{T}
    θ3::T
end
(p::Poly3)(x::Real) = evalpoly(x, p.θ3)

poly3s = Poly3([10, 1, 0.1])
# poly3s is now a valid Flux model

To resume training after an interruption, you should save both the model state and the optimizer state. This allows you to restart training exactly where you left off.

Use Flux.state(model) for the model parameters and the state object returned by Flux.setup for the optimizer.

using Flux
using JLD2

model = Chain(Dense(10 => 5, relu), Dense(5 => 2))
opt_state = Flux.setup(AdamW(), model)

# ... during training loop ...

# Save both to a checkpoint file
model_state = Flux.state(model)
jldsave("checkpoint_epoch=42.jld2"; model_state, opt_state)

While you can compose functions using the Julia composition operator ∘, Flux provides Flux.Chain for a more structured approach.

Key differences between Chain and standard function composition:

• Chain works left-to-right (the reverse of ∘).
• Layers are stored in a tuple, allowing access via model.layers[i].weight.
• Built-in layers like Dense use glorot_uniform initialization by default and include performance optimizations for batch processing and memory reuse.

Many normalisation layers (like BatchNorm) behave differently during training versus inference (testing).

Flux attempts to automatically detect the mode, which works best with Zygote. If you are using other automatic differentiation packages (like Tracker, Yota, or ForwardDiff), automatic detection may fail.

You can manually switch the mode for an entire model using Flux.trainmode! and Flux.testmode!.

To train neural networks, you typically use DataLoader to handle minibatch learning. This involves splitting your data into training and validation sets and wrapping them in a DataLoader with a specified batchsize and shuffle preference. For classification tasks, use onehotbatch to convert labels into a format suitable for loss functions like logitcrossentropy.

using Flux
using MLUtils: splitobs, numobs
using MLDatasets: CIFAR10

# Load data
train_x, train_y = CIFAR10(:train)[:]
labels = onehotbatch(train_y, 0:9)

# Split into training and validation sets
trainset, valset = splitobs((train_x, labels), at = 45000)

# Create DataLoaders
trainloader = DataLoader(trainset, batchsize = 1000, shuffle = true)
valloader = DataLoader(trainset, batchsize = 1000)