Machine Learning with Flux
Prerequisites: Optimization
Machine-learning models are functions with parameters, trained by gradient-based optimization — exactly what Julia is good at.
Logistic regression from scratch
using Random, LinearAlgebra, Statistics
rng = Xoshiro(7)
n = 200
X = randn(rng, 2, n) # 2 features × 200 samples
y = Float64.(X[1, :] .+ 2 .* X[2, :] .> 0) # true decision boundary
σ(z) = 1 / (1 + exp(-z))
function train(X, y; lr = 0.5, epochs = 500)
w = zeros(size(X, 1)); b = 0.0
for _ in 1:epochs
p = σ.(w' * X .+ b) # predictions, 1 × n
err = vec(p) .- y
w .-= lr .* (X * err) ./ length(y)
b -= lr * mean(err)
end
return w, b
end
w, b = train(X, y)
accuracy = mean((vec(σ.(w' * X .+ b)) .> 0.5) .== (y .== 1))
println("accuracy: ", accuracy)
println("weight ratio w2/w1 ≈ ", round(w[2] / w[1], digits = 1))
accuracy: 0.995
weight ratio w2/w1 ≈ 2.1
The learned weights point in the direction of the true boundary
(x1 + 2 x2 = 0, a ratio of about 2). The gradient here was derived by
hand; for anything bigger, let a library do it.
Flux.jl
Flux.jl is Julia's most popular deep-learning library. Models are ordinary Julia structs, and Zygote.jl computes gradients of ordinary Julia code:
# illustrative — needs the Flux package
using Flux
model = Chain(
Dense(2 => 16, relu),
Dense(16 => 16, relu),
Dense(16 => 1, sigmoid),
)
X = randn(Float32, 2, 1000)
y = reshape(Float32.(X[1, :] .^ 2 .+ X[2, :] .^ 2 .< 1), 1, :) # inside a circle?
loss(m, x, y) = Flux.binarycrossentropy(m(x), y)
opt_state = Flux.setup(Adam(0.01), model)
data = Flux.DataLoader((X, y), batchsize = 64, shuffle = true)
for epoch in 1:200
Flux.train!(loss, model, data, opt_state)
end
accuracy = mean((model(X) .> 0.5) .== (y .> 0.5))
Dense(in => out, activation)layers are combined withChain.Flux.setupcreates optimizer state;Flux.train!loops over the data, computes gradients, and updates the parameters.- Move the model and data to a GPU with
gpu(model)(with CUDA.jl, AMDGPU.jl, or Metal.jl loaded).
A training loop written out explicitly:
# illustrative — needs the Flux package
for (x, y) in data
grads = Flux.gradient(m -> loss(m, x, y), model)
Flux.update!(opt_state, model, grads[1])
end
The wider ML ecosystem
- Lux.jl — explicit-parameter neural networks, popular in scientific ML.
- MLJ.jl — a unified interface to classical models (trees, SVMs, clustering), similar to scikit-learn.
- Enzyme.jl — fast automatic differentiation at the LLVM level.
- Transformers.jl, Metalhead.jl (vision models), and MLDatasets.jl.
Julia's niche in ML is research and scientific machine learning — custom models, physics-informed networks, and differentiable simulations — rather than training the largest models, where Python frameworks dominate.