Skip to content
AgileLogics
7 min read

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 with Chain.
  • Flux.setup creates 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.