DOCSEARCHbyAlgolia
Submit Docs

Lux.jl Documentation

website·Indexed 19 days ago

https://lux.csail.mit.edu/dev

Documentation for Lux.jl, a neural network library for Julia. It includes an introduction to the library, migration guides for Lux v1, and a comprehensive set of tutorials covering MLP polynomial fitting, LSTMs, HyperNetworks, Physics-Informed Neural Networks (PINNs), Convolutional VAEs, Graph Convolutional Networks (GCN), and Neural ODEs for gravitational waveforms.

Tokens
76.9K
Snippets
439
Records
506
Agent score
99%

What's inside Lux.jl

  2. Overview of LuxTestUtils

    LuxTestUtils is a testing package providing utilities for verifying gradient correctness and dynamic dispatch of Lux.jl models.

    Important Usage Note: This package is intended exclusively for testing. To avoid increasing load times in production, it is recommended not to add LuxTestUtils as a dependency in your main package's Project.toml.

  3. Overview of Lux.jl Design Principles

    Lux is a neural network framework for Julia built using pure functions to be compiler and autodiff friendly. Unlike some other frameworks, Lux avoids coupled models and parameters and internal mutations. Its core design principles are:

    • Immutability: Layers are immutable and cannot store parameters or state; they only store the information needed to construct them.
    • Pure Functions: Layers are implemented as pure functions.
    • State Management: Layers return a Tuple containing the result and the updated state.
    • Determinism: Given the same inputs, outputs must be identical. Stochastic functions (e.g., Lux.Dropout) must control randomness using rngs passed within the state.
    • Extensibility: The framework is designed to be easily extensible and is extensively tested across various AD backends and hardware.

  5. Train an LSTM Encoder-Decoder model with Lux.jl

    To train an LSTM Encoder-Decoder model, use a training loop that leverages Training.single_train_step! with an optimizer (e.g., Adam) and a loss function (e.g., MSELoss). The model architecture typically consists of an RNNEncoder and an RNNDecoder using LSTMCell. The training process can incorporate 'mixed teacher forcing' to improve convergence by blending ground-truth targets and model predictions during the decoding phase.
    function train(
    train_dataset,
    validation_dataset;
    nepochs=50,
    batchsize=32,
    hidden_dims=32,
    training_mode=:mixed_teacher_forcing,
    teacher_forcing_ratio=0.5f0,
    learning_rate=1e-3,
)
    (X_train, Y_train), (X_test, Y_test) = train_dataset, validation_dataset
    in_dims = size(X_train, 1)
    target_len = size(Y_train, 2)

    train_dataloader = DataLoader(
        (X_train, Y_train);
        batchsize=min(batchsize, size(X_train, 4)),
        shuffle=true,
        partial=false,
    ) |> xdev
    X_test, Y_test = (X_test, Y_test) |> xdev

    model = RNNEncoderDecoder(
        RNNEncoder(LSTMCell(in_dims => hidden_dims)),
        RNNDecoder(
            LSTMCell(in_dims => hidden_dims),
            Dense(hidden_dims => in_dims);
            training_mode,
            teacher_forcing_ratio,
        ),
    )
    ps, st = Lux.setup(Random.default_rng(), model) |> xdev

    train_state = Training.TrainState(model, ps, st, Optimisers.Adam(learning_rate))

    for epoch in 1:nepochs
        for (x, y) in train_dataloader
            (_, _, _, train_state) = Training.single_train_step!(
                AutoEnzyme(),
                MSELoss(),
                ((x, target_len, y), y),
                train_state;
                return_gradients=Val(false),
            )
        end
        # Validation logic here...
    end

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

  7. Implement a ResNet-20 architecture using Lux.jl

    Build a ResNet-20 model for image classification (e.g., CIFAR-10) using Lux.jl. The architecture consists of an initial convolutional layer, three stages of residual blocks (with increasing channel depths and stride-2 downsampling), global mean pooling, and a final dense layer for classification.
    function ConvBN(kernel_size, (in_chs, out_chs), act; kwargs...)
    return Chain(Conv(kernel_size, in_chs => out_chs, act; kwargs...), BatchNorm(out_chs))
end

function BasicBlock(in_channels, out_channels; stride=1)
    connection = if (stride == 1 && in_channels == out_channels)
        NoOpLayer()
    else
        Conv((3, 3), in_channels => out_channels, identity; stride=stride, pad=SamePad())
    end
    return Chain(
        Parallel(
            +,
            connection,
            Chain(
                ConvBN((3, 3), in_channels => out_channels, relu; stride, pad=SamePad()),
                ConvBN((3, 3), out_channels => out_channels, identity; pad=SamePad()),
            ),
        ),
        Base.BroadcastFunction(relu),
    )
end

function ResNet20(; num_classes=10)
    layers = []

    # Initial Conv Layer
    push!(layers, Chain(Conv((3, 3), 3 => 16, relu; pad=SamePad()), BatchNorm(16)))

    # Residual Blocks
    block_configs = [
        # (in_channels, out_channels, num_blocks, stride)
        (16, 16, 3, 1),
        (16, 32, 3, 2),
        (32, 64, 3, 2),
    ]

    for (in_channels, out_channels, num_blocks, stride) in block_configs
        for i in 1:num_blocks
            push!(
                layers,
                BasicBlock(
                    i == 1 ? in_channels : out_channels,
                    out_channels;
                    stride=(i == 1 ? stride : 1),
                ),
            )
        end
    end

    # Global Pooling and Final Dense Layer
    push!(layers, GlobalMeanPool())
    push!(layers, FlattenLayer())
    push!(layers, Dense(64 => num_classes))

    return Chain(layers...)
end

  8. Configure DDIM model hyperparameters

    The DDIM model architecture and signal processing are controlled by the following hyperparameters, which can be passed to the main function or configured via CLI arguments:
    # Hyperparameter defaults
channels = [32, 64, 96, 128] # UNet channels per stage
block_depth = 2              # Number of residual blocks per stage
min_freq = 1.0f0             # Sinusoidal embedding min frequency
max_freq = 1000.0f0          # Sinusoidal embedding max frequency
embedding_dims = 32          # Sinusoidal embedding dimension
min_signal_rate = 0.02f0     # Minimum signal rate
max_signal_rate = 0.95f0     # Maximum signal rate

  9. Implement custom layers in Lux

    Unlike Flux, which stores parameters within the layer struct, Lux layers are stateless descriptions of the architecture.

    Key implementation rules:

    1. Avoid Mutables: Do not store Arrays or other mutable structures directly inside a Lux Layer struct. Use lazy initialization (functions) instead.
    2. Parameters vs. States: Use Lux.initialparameters for trainable weights and Lux.initialstates for non-trainable state.
    3. Device Transfer: Because layers are stateless, device transfer utilities (like gpu_device) must be applied to the parameters and states, not the layer itself.
    using Lux, Random, NNlib, Zygote

struct LuxLinear <: Lux.AbstractLuxLayer
    init_A
    init_B
end

# Constructor using lazy initialization to avoid storing arrays in the struct
function LuxLinear(A::AbstractArray, B::AbstractArray)
    return LuxLinear(() -> copy(A), () -> copy(B))
end

# Define trainable parameters
Lux.initialparameters(::AbstractRNG, layer::LuxLinear) = (B=layer.init_B(),)

# Define non-trainable state
Lux.initialstates(::AbstractRNG, layer::LuxLinear) = (A=layer.init_A(),)

# Forward pass: returns output and updated state
(l::LuxLinear)(x, ps, st) = st.A * ps.B * x, st

  10. Understand LuxCore layer use cases

    LuxCore provides abstract layers for Lux and is primarily used in the following scenarios:

    • Internal implementation of the @compact macro.
    • SciML codebases where state propagation requires boxing to avoid type instability (e.g., in Neural ODEs).
    • Enabling Nested Automatic Differentiation (AD) support within Lux.

  11. Implement a PINN architecture for 2D PDEs in Lux.jl

    To solve 2D PDEs, a neural network is required that takes 3 input variables (typically x, y, and t) and outputs a scalar value. The following implementation uses a Chain of Dense layers with tanh activations, wrapped in a custom PINN struct to facilitate training with Training.TrainState.
    struct PINN{M} <: AbstractLuxWrapperLayer{:model}
    model::M
end

function PINN(; hidden_dims::Int=32)
    return PINN(
        Chain(
            Dense(3 => hidden_dims, tanh),
            Dense(hidden_dims => hidden_dims, tanh),
            Dense(hidden_dims => hidden_dims, tanh),
            Dense(hidden_dims => 1),
        ),
    )
end

  12. Use tanh_fast for high-performance tanh activation

    tanh_fast(x) is a faster but slightly less accurate version of tanh. For Float32 inputs, it is typically about 10 times faster than the standard tanh function. The error is approximately 5 eps (compared to under 2 eps for standard tanh). For types other than Float32 or Float64, it defaults to calling the standard tanh function.
    julia> tanh(0.5f0)
0.46211717f0

julia> tanh_fast(0.5f0)
0.46211714f0