Flux is an elegant approach to machine learning. It's a 100% pure-Julia stack, and provides lightweight abstractions on top of Julia's native GPU and AD support. Flux makes the easy things easy while remaining fully hackable.
- Layers and losses for MLPs, convolutional, recurrent and attention models, with normalisation, pooling, embeddings and more.
- Your choice of AD engine: Zygote by default, or Enzyme and Mooncake, see gradients.
- GPU support for NVIDIA, AMD and Apple hardware, see the GPU guide.
- Compilation with Reactant: the whole training step fused into one XLA program, on CPU, GPU or TPU.
- Distributed training with MPI and NCCL, see the GPU guide.
Works best with Julia 1.10 or later. Install Flux from the Julia REPL with using Pkg; Pkg.add("Flux"), then try this very short example:
using Flux
# XOR problem: classify the points of the unit square by quadrant
noisy = rand(Float32, 2, 1000) # one point per column
truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]
target = Flux.onehotbatch(truth, [true, false])
loader = Flux.DataLoader((noisy, target), batchsize=64, shuffle=true)
model = Chain(Dense(2 => 3, tanh), BatchNorm(3), Dense(3 => 2))
opt_state = Flux.setup(Adam(0.01), model)
for epoch in 1:1_000
Flux.train!(model, loader, opt_state) do m, x, y
Flux.logitcrossentropy(m(x), y)
end
end
accuracy = sum(Flux.onecold(model(noisy), [true, false]) .== truth) / length(truth)
To train on a GPU, load the package for your hardware and move the model and the data loader to the device:
using CUDA # or AMDGPU, Metal
device = gpu_device() # the CPU, if no GPU is available
model = model |> device # before calling Flux.setup
loader = loader |> device # moves one batch at a time
Flux models are plain Julia. In fact, almost any parameterised function is a valid model, such as the closure over w, b, v below. The same function can also be implemented with built-in layers as shown.
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:300
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")
The quickstart page has a longer version of the first example, with plots. See the documentation for details, the model zoo for examples, or the examples folder for larger end-to-end scripts: GPT-2 pretraining and ResNet on TinyImageNet. Ask questions on the Julia discourse or slack.
If you use Flux in your research, please cite our work.
