Flux.jl: machine learning en Julia con diferenciación automática nativa

PyTorch y TensorFlow dominan el espacio del machine learning, y probablemente van a seguir haciéndolo. Pero tienen una particularidad interesante: sus operaciones rápidas están escritas en C++ y CUDA, y la capa de Python que el usuario toca es más o menos un frontend para esas bibliotecas compiladas. Flux.jl toma otro camino: el framework está escrito en Julia, y la diferenciación automática también.

Zygote.jl: diferenciación automática en Julia puro

La clave de cualquier framework de ML es la diferenciación automática: la capacidad de calcular gradientes de funciones arbitrarias. Zygote.jl lo hace en modo reverse (backward pass), que es el eficiente para redes neuronales con muchos parámetros y pocas salidas.

using Zygote

# Gradiente de una función escalar
f(x) = x^3 + 2x^2 - x + 1
gradient(f, 3.0)   # devuelve (28.0,) — df/dx en x=3

# Gradiente de una función multivariable
g(x, y) = x^2 + x*y + y^2
gradient(g, 2.0, 3.0)  # devuelve (7.0, 7.0) — gradiente en (2,3)

# Gradient con arrays
h(w) = sum(w .^ 2)
gradient(h, [1.0, 2.0, 3.0])  # devuelve ([2.0, 4.0, 6.0],)

Lo que hace Zygote diferente es que puede diferenciar código Julia arbitrario: bucles, condicionales, llamadas a otras funciones. No estás limitado a un conjunto de operaciones predefinidas como en los grafos estáticos de TensorFlow 1.x.

Flux.jl: construir redes neuronales

Chain es el constructor principal para redes secuenciales. Encadena capas y el resultado es un modelo que aplica cada capa en orden:

using Flux

# Red neuronal simple para clasificación
model = Chain(
    Dense(784, 256, relu),   # 784 entradas, 256 neuronas, activación ReLU
    Dropout(0.2),             # dropout del 20% para regularización
    Dense(256, 128, relu),
    Dense(128, 10),           # 10 clases de salida
    softmax                   # probabilidades de clase
)

# Ver parámetros del modelo
println("Parámetros totales: ", sum(length, Flux.params(model)))

# Forward pass
x = rand(Float32, 784, 32)  # batch de 32 imágenes
y_pred = model(x)            # predicciones (10, 32)

Las capas disponibles cubren los casos habituales: Dense, Conv, MaxPool, RNN, LSTM, GRU, Embedding, LayerNorm, BatchNorm, y Dropout. Para arquitecturas no estándar puedes definir tus propias capas como structs de Julia.

Entrenamiento

using Flux, Flux.Losses

# Datos de ejemplo (MNIST simplificado)
x_train = rand(Float32, 784, 1000)   # 1000 ejemplos
y_train = Flux.onehotbatch(rand(1:10, 1000), 1:10)  # one-hot

# Función de pérdida
loss(x, y) = Flux.crossentropy(model(x), y)

# Optimizador
opt = ADAM(0.001)

# Parámetros del modelo
ps = Flux.params(model)

# Bucle de entrenamiento
data = [(x_train[:, i:i+31], y_train[:, i:i+31]) for i in 1:32:969]

for epoch in 1:10
    Flux.train!(loss, ps, data, opt)
    println("Epoch $epoch, loss: ", loss(x_train, y_train))
end

Capas convolucionales para imágenes

# CNN para clasificación de imágenes
cnn = Chain(
    Conv((3, 3), 1 => 32, relu, pad=1),   # 1 canal -> 32 filtros, kernel 3x3
    MaxPool((2, 2)),
    Conv((3, 3), 32 => 64, relu, pad=1),
    MaxPool((2, 2)),
    Flux.flatten,                           # aplanar para la capa densa
    Dense(64 * 7 * 7, 128, relu),
    Dense(128, 10),
    softmax
)

# Entrada: (ancho, alto, canales, batch) — orden WHCN en Flux
x_img = rand(Float32, 28, 28, 1, 16)  # batch de 16 imágenes 28x28 en escala de grises
output = cnn(x_img)

Transferencia al GPU

Mover el modelo y los datos al GPU es una línea:

using CUDA

# Mover modelo a GPU
model_gpu = model |> gpu

# Mover datos a GPU
x_gpu = x_train |> gpu
y_gpu = y_train |> gpu

# Entrenar en GPU (mismo código)
ps_gpu = Flux.params(model_gpu)
Flux.train!(
    (x, y) -> Flux.crossentropy(model_gpu(x), y),
    ps_gpu, data_gpu, opt
)

Definir capas personalizadas

Una de las ventajas de que Flux esté en Julia puro es que crear capas nuevas no requiere C++. Es un struct con una función de forward pass:

# Capa personalizada: atención simple
struct Atencion
    W::Matrix{Float32}
end

# Constructor
Atencion(d_in, d_out) = Atencion(randn(Float32, d_out, d_in) * 0.01f0)

# Forward pass
function (a::Atencion)(x)
    scores = a.W * x
    return softmax(scores)
end

# Registrar parámetros para que Flux los optimice
Flux.@functor Atencion

# Usar en un Chain como cualquier otra capa
model_custom = Chain(Dense(128, 64, relu), Atencion(64, 10))

Esta integración sin fricciones entre capas de usuario y capas de librería viene del sistema de tipos de Julia, que ya explicamos en el artículo sobre multiple dispatch. Flux no necesita saber nada sobre tu capa: solo necesita que tenga parámetros registrados y una función forward.

Comparado con Python 3.13 con su nuevo JIT, el enfoque de Julia tiene la ventaja de que la diferenciación automática puede atravesar cualquier código Julia, incluyendo solvers de ecuaciones diferenciales. Esto permite cosas como diferenciar a través de una simulación física, que es difícil de hacer en PyTorch sin implementaciones especiales.

Imagen: Pexels / Google DeepMind