### Old Tutorials/2021 02 07 Convnet
# [Tutorial: A Simple ConvNet](@id man-convnet-tutorial)
In this tutorial, we build a simple Convolutional Neural Network (ConvNet) to classify the MNIST dataset. This model has a simple architecture with three feature detection layers (Conv -> ReLU -> MaxPool) followed by a final dense layer that classifies MNIST handwritten digits. Note that this model, while simple, should hit around 99% test accuracy after training for approximately 20 epochs.
This example writes out the saved model to the file `mnist_conv.bson`. Also, it demonstrates basic model construction, training, saving, conditional early-exit, and learning rate scheduling.
To run this example, we need the following packages:
```julia
using Flux, MLDatasets, Statistics
using Flux: onehotbatch, onecold, logitcrossentropy, params
using MLDatasets: MNIST
using Base.Iterators: partition
using Printf, BSON
using CUDA
CUDA.allowscalar(false)
```
We set default values for learning rate, batch size, number of epochs, and path for saving the file `mnist_conv.bson`:
```julia
Base.@kwdef mutable struct TrainArgs
lr::Float64 = 3e-3
epochs::Int = 20
batch_size = 128
savepath::String = "./"
end
```
## Data
To train our model, we need to bundle images together with their labels and group them into mini-batches (makes the training process faster). We define the function `make_minibatch` that takes as inputs the images (`X`) and their labels (`Y`) as well as the indices for the mini-batches (`idx`):
```julia
function make_minibatch(X, Y, idxs)
X_batch = Array{Float32}(undef, size(X)[1:end-1]..., 1, length(idxs))
for i in 1:length(idxs)
X_batch[:, :, :, i] = Float32.(X[:,:,idxs[i]])
end
Y_batch = onehotbatch(Y[idxs], 0:9)
return (X_batch, Y_batch)
end
```
`make_minibatch` takes the following steps:
* Creates the `X_batch` array of size `28x28x1x128` to store the mini-batches.
* Stores the mini-batches in `X_batch`.
* One hot encodes the labels of the images.
* Stores the labels in `Y_batch`.
`get_processed_data` loads the train and test data from `Flux.Data.MNIST`. First, it loads the images and labels of the train data set, and creates an array that contains the indices of the train images that correspond to each mini-batch (of size `args.batch_size`). Then, it calls the `make_minibatch` function to create all of the train mini-batches. Finally, it loads the test images and creates one mini-batch that contains them all.
```julia
function get_processed_data(args)
# Load labels and images
train_imgs, train_labels = MNIST.traindata()
mb_idxs = partition(1:length(train_labels), args.batch_size)
train_set = [make_minibatch(train_imgs, train_labels, i) for i in mb_idxs]
# Prepare test set as one giant minibatch:
test_imgs, test_labels = MNIST.testdata()
test_set = make_minibatch(test_imgs, test_labels, 1:length(test_labels))
return train_set, test_set
end
```
## Model
Now, we define the `build_model` function that creates a ConvNet model which is composed of *three* convolution layers (feature detection) and *one* classification layer. The input layer size is `28x28`. The images are grayscale, which means there is only *one* channel (compared to 3 for RGB) in every data point. Combined together, the convolutional layer structure would look like `Conv(kernel, input_channels => output_channels, ...)`. Each convolution layer reduces the size of the image by applying the Rectified Linear unit (ReLU) and MaxPool operations.
On the other hand, the classification layer outputs a vector of 10 dimensions (a dense layer), that is, the number of classes that the model will be able to predict.
```julia
function build_model(args; imgsize = (28,28,1), nclasses = 10)
cnn_output_size = Int.(floor.([imgsize[1]/8,imgsize[2]/8,32]))
return Chain(
# First convolution, operating upon a 28x28 image
Conv((3, 3), imgsize[3]=>16, pad=(1,1), relu),
MaxPool((2,2)),
# Second convolution, operating upon a 14x14 image
Conv((3, 3), 16=>32, pad=(1,1), relu),
MaxPool((2,2)),
# Third convolution, operating upon a 7x7 image
Conv((3, 3), 32=>32, pad=(1,1), relu),
MaxPool((2,2)),
# Reshape 3d array into a 2d one using `Flux.flatten`, at this point it should be (3, 3, 32, N)
flatten,
Dense(prod(cnn_output_size), 10))
end
```
To chain the layers of a model we use the Flux function [Chain](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Chain). It enables us to call the layers in sequence on a given input. Also, we use the function [flatten](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.flatten) to reshape the output image from the last convolution layer. Finally, we call the [Dense](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Dense) function to create the classification layer.
## Training
Before training our model, we need to define a few functions that will be helpful for the process:
* `augment` adds gaussian random noise to our image, to make it more robust:
* `anynan` checks whether any element of the params is NaN or not:
* `accuracy` computes the proportion of inputs `x` correctly classified by our ConvNet:
```julia
augment(x) = x .+ gpu(0.1f0*randn(eltype(x), size(x)))
anynan(x) = any(y -> any(isnan, y), x)
accuracy(x, y, model) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))
```
Finally, we define the `train` function:
```
/* Detailed source-code truncated for AI context efficiency. */
```
`train` calls the functions we defined above and trains our model. It stops when the model achieves 99% accuracy (early-exiting) or after performing 20 steps. More specifically, it performs the following steps:
* Loads the MNIST dataset.
* Builds our ConvNet model (as described above).
* Loads the train and test data sets as well as our model onto a GPU (if available).
* Defines a `loss` function that calculates the crossentropy between our prediction and the ground truth.
* Sets the [Adam optimiser](https://fluxml.ai/Flux.jl/stable/training/optimisers/#Flux.Optimise.Adam) to train the model with learning rate `args.lr`.
* Runs the training loop. For each step (or epoch), it executes the following:
* Calls `Flux.train!` function to execute one training step.
* If any of the parameters of our model is `NaN`, then the training process is terminated.
* Calculates the model accuracy.
* If the model accuracy is >= 0.999, then early-exiting is executed.
* If the actual accuracy is the best so far, then the model is saved to `mnist_conv.bson`. Also, the new best accuracy and the current epoch is saved.
* If there has not been any improvement for the last 5 epochs, then the learning rate is dropped and the process waits a little longer for the accuracy to improve.
* If the last improvement was more than 10 epochs ago, then the process is terminated.
## Testing
Finally, to test our model we define the `test` function:
```julia
function test(; kws...)
args = TrainArgs(; kws...)
# Loading the test data
_,test_set = get_processed_data(args)
# Re-constructing the model with random initial weights
model = build_model(args)
# Loading the saved parameters
BSON.@load joinpath(args.savepath, "mnist_conv.bson") params
# Loading parameters onto the model
Flux.loadparams!(model, params)
test_set = gpu.(test_set)
model = gpu(model)
@show accuracy(test_set...,model)
end
```
`test` loads the MNIST test data set, reconstructs the model, and loads the saved parameters (in `mnist_conv.bson`) onto it. Finally, it computes our model's predictions for the test set and shows the test accuracy (around 99%).
To see the full version of this example, see [Simple ConvNets - model-zoo](https://github.com/FluxML/model-zoo/blob/master/vision/conv_mnist/conv_mnist.jl).
## Resources
* [Neural Networks in Flux.jl with Huda Nassar (working with the MNIST dataset)](https://youtu.be/Oxi0Pfmskus)
* [Convolutional Neural Networks (CNNs / ConvNets)](https://cs231n.github.io/convolutional-networks/).
* [Convolutional Neural Networks Tutorial in PyTorch](https://adventuresinmachinelearning.com/convolutional-neural-networks-tutorial-in-pytorch/).
!!! info
Originally published at [fluxml.ai](https://fluxml.ai/tutorials/) on 7 February 2021.
Written by Elliot Saba, Adarsh Kumar, Mike J Innes, Dhairya Gandhi, Sudhanshu Agrawal, Sambit Kumar Dash, fps.io, Carlo Lucibello, Andrew Dinhobl, Liliana Badillo
---
### Old Tutorials/2021 10 08 Dcgan Mnist
# Deep Convolutional Generative Adversarial Network (DCGAN)
This is a beginner level tutorial for generating images of handwritten digits using a [Deep Convolutional Generative Adversarial Network](https://arxiv.org/pdf/1511.06434.pdf) inspired by the [TensorFlow tutorial on DCGAN](https://www.tensorflow.org/tutorials/generative/dcgan).
## What are GANs?
[Generative Adversarial Neural Networks or simply GANs](https://arxiv.org/abs/1406.2661) introduced by Goodfellow et al. is one of the most innovative ideas in modern-day machine learning. GANs are used extensively in the field of image and audio processing to generate high-quality synthetic data that can easily be passed off as real data.
A GAN is composed of two sub-models - the **generator** and the **discriminator** acting against one another. The generator can be considered as an artist who draws (generates) new images that look real, whereas the discriminator is a critic who learns to tell real images apart from fakes.
The GAN starts with a generator and discriminator which have very little or no idea about the underlying data. During training, the generator progressively becomes better at creating images that look real, while the discriminator becomes better at telling them apart. The process reaches equilibrium when the discriminator can no longer distinguish real images from fakes.
[[source]](https://www.tensorflow.org/tutorials/generative/dcgan)
This tutorial demonstrates the process of training a DC-GAN on the [MNIST dataset for handwritten digits](http://yann.lecun.com/exdb/mnist/). The following animation shows a series of images produced by the generator as it was trained for 25 epochs. The images begin as random noise, but over time, the images become increasingly similar to handwritten numbers.
## Setup
We need to install some Julia packages before we start with our implementation of DCGAN.
```julia
using Pkg
# Activate a new project environment in the current directory
Pkg.activate(".")
# Add the required packages to the environment
Pkg.add(["Images", "Flux", "MLDatasets", "CUDA", "Parameters"])
```
*Note: Depending on your internet speed, it may take a few minutes for the packages install.*
After installing the libraries, load the required packages and functions:
```julia
using Base.Iterators: partition
using Printf
using Statistics
using Random
using Images
using Flux: params, DataLoader
using Flux.Optimise: update!
using Flux.Losses: logitbinarycrossentropy
using MLDatasets: MNIST
using CUDA
```
Now we set default values for the learning rates, batch size, epochs, the usage of a GPU (if available) and other hyperparameters for our model.
```julia
Base.@kwdef struct HyperParams
batch_size::Int = 128
latent_dim::Int = 100
epochs::Int = 25
verbose_freq::Int = 1000
output_dim::Int = 5
disc_lr::Float64 = 0.0002
gen_lr::Float64 = 0.0002
device::Function = gpu
end
```
## Loading the data
As mentioned before, we will be using the MNIST dataset for handwritten digits. So we begin with a simple function for loading and pre-processing the MNIST images:
```julia
function load_MNIST_images(hparams)
images = MNIST.traintensor(Float32)
# Normalize the images to (-1, 1)
normalized_images = @. 2f0 * images - 1f0
image_tensor = reshape(normalized_images, 28, 28, 1, :)
# Create a dataloader that iterates over mini-batches of the image tensor
dataloader = DataLoader(image_tensor, batchsize=hparams.batch_size, shuffle=true)
return dataloader
end
```
To learn more about loading images in Flux, you can check out [this tutorial](https://fluxml.ai/tutorials/2021/01/21/data-loader.html).
*Note: The data returned from the dataloader is loaded is on the CPU. To train on the GPU, we need to transfer the data to the GPU beforehand.*
## Create the models
### Generator
Our generator, a.k.a. the artist, is a neural network that maps low dimensional data to a high dimensional form.
- This low dimensional data (seed) is generally a vector of random values sampled from a normal distribution.
- The high dimensional data is the generated image.
The `Dense` layer is used for taking the seed as an input which is upsampled several times using the [ConvTranspose](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.ConvTranspose) layer until we reach the desired output size (in our case, 28x28x1). Furthermore, after each `ConvTranspose` layer, we apply the Batch Normalization to stabilize the learning process.
We will be using the [relu](https://fluxml.ai/Flux.jl/stable/models/nnlib/#NNlib.relu) activation function for each layer except the output layer, where we use `tanh` activation.
We will also apply the weight initialization method mentioned in the original DCGAN paper.
```julia
# Function for initializing the model weights with values
# sampled from a Gaussian distribution with μ=0 and σ=0.02
dcgan_init(shape...) = randn(Float32, shape) * 0.02f0
```
```julia
function Generator(latent_dim)
Chain(
Dense(latent_dim => 7*7*256, bias=false),
BatchNorm(7*7*256, relu),
x -> reshape(x, 7, 7, 256, :),
ConvTranspose((5, 5), 256 => 128; stride = 1, pad = 2, init = dcgan_init, bias=false),
BatchNorm(128, relu),
ConvTranspose((4, 4), 128 => 64; stride = 2, pad = 1, init = dcgan_init, bias=false),
BatchNorm(64, relu),
# The tanh activation ensures that output is in range of (-1, 1)
ConvTranspose((4, 4), 64 => 1, tanh; stride = 2, pad = 1, init = dcgan_init, bias=false),
)
end
```
Time for a small test!! We create a dummy generator and feed a random vector as a seed to the generator. If our generator is initialized correctly it will return an array of size (28, 28, 1, `batch_size`). The `@assert` macro in Julia will raise an exception for the wrong output size.
```julia
# Create a dummy generator of latent dim 100
generator = Generator(100)
noise = randn(Float32, 100, 3) # The last axis is the batch size
# Feed the random noise to the generator
gen_image = generator(noise)
@assert size(gen_image) == (28, 28, 1, 3)
```
Our generator model is yet to learn the correct weights, so it does not produce a recognizable image for now. To train our poor generator we need its equal rival, the *discriminator*.
### Discriminator
The Discriminator is a simple CNN based image classifier. The `Conv` layer a is used with a [leakyrelu](https://fluxml.ai/Flux.jl/stable/models/nnlib/#NNlib.leakyrelu) activation function.
```julia
function Discriminator()
Chain(
Conv((4, 4), 1 => 64; stride = 2, pad = 1, init = dcgan_init),
x->leakyrelu.(x, 0.2f0),
Dropout(0.3),
Conv((4, 4), 64 => 128; stride = 2, pad = 1, init = dcgan_init),
x->leakyrelu.(x, 0.2f0),
Dropout(0.3),
# The output is now of the shape (7, 7, 128, batch_size)
flatten,
Dense(7 * 7 * 128, 1)
)
end
```
For a more detailed implementation of a CNN-based image classifier, you can refer to [this tutorial](https://fluxml.ai/tutorials/2021/02/07/convnet.html).
Now let us check if our discriminator is working:
```julia
# Dummy Discriminator
discriminator = Discriminator()
# We pass the generated image to the discriminator
logits = discriminator(gen_image)
@assert size(logits) == (1, 3)
```
Just like our dummy generator, the untrained discriminator has no idea about what is a real or fake image. It needs to be trained alongside the generator to output positive values for real images, and negative values for fake images.
## Loss functions for GAN
In a GAN problem, there are only two labels involved: fake and real. So Binary CrossEntropy is an easy choice for a preliminary loss function.
But even if Flux's `binarycrossentropy` does the job for us, due to numerical stability it is always preferred to compute cross-entropy using logits. Flux provides [logitbinarycrossentropy](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.logitbinarycrossentropy) specifically for this purpose. Mathematically it is equivalent to `binarycrossentropy(σ(ŷ), y, kwargs...).`
### Discriminator Loss
The discriminator loss quantifies how well the discriminator can distinguish real images from fakes. It compares
- discriminator's predictions on real images to an array of 1s, and
- discriminator's predictions on fake (generated) images to an array of 0s.
These two losses are summed together to give a scalar loss. So we can write the loss function of the discriminator as:
```julia
function discriminator_loss(real_output, fake_output)
real_loss = logitbinarycrossentropy(real_output, 1)
fake_loss = logitbinarycrossentropy(fake_output, 0)
return real_loss + fake_loss
end
```
### Generator Loss
The generator's loss quantifies how well it was able to trick the discriminator. Intuitively, if the generator is performing well, the discriminator will classify the fake images as real (or 1).
```julia
generator_loss(fake_output) = logitbinarycrossentropy(fake_output, 1)
```
We also need optimisers for our network. Why you may ask? Read more [here](https://towardsdatascience.com/overview-of-various-optimisers-in-neural-networks-17c1be2df6d5). For both the generator and discriminator, we will use the [ADAM optimiser](https://fluxml.ai/Flux.jl/stable/training/optimisers/#Flux.Optimise.ADAM).
## Utility functions
The output of the generator ranges from (-1, 1), so it needs to be de-normalized before we can display it as an image. To make things a bit easier, we define a function to visualize the output of the generator as a grid of images.
```julia
function create_output_image(gen, fixed_noise, hparams)
fake_images = cpu(gen.(fixed_noise))
image_array = reduce(vcat, reduce.(hcat, partition(fake_images, hparams.output_dim)))
image_array = permutedims(dropdims(image_array; dims=(3, 4)), (2, 1))
image_array = @. Gray(image_array + 1f0) / 2f0
return image_array
end
```
## Training
For the sake of simplifying our training problem, we will divide the generator and discriminator training into two separate functions.
```julia
function train_discriminator!(gen, disc, real_img, fake_img, opt, ps, hparams)
disc_loss, grads = Flux.withgradient(ps) do
discriminator_loss(disc(real_img), disc(fake_img))
end
# Update the discriminator parameters
update!(opt, ps, grads)
return disc_loss
end
```
We define a similar function for the generator.
```julia
function train_generator!(gen, disc, fake_img, opt, ps, hparams)
gen_loss, grads = Flux.withgradient(ps) do
generator_loss(disc(fake_img))
end
update!(opt, ps, grads)
return gen_loss
end
```
Now that we have defined every function we need, we integrate everything into a single `train` function where we first set up all the models and optimisers and then train the GAN for a specified number of epochs.
```
/* Detailed source-code truncated for AI context efficiency. */
```
Now we finally get to train the GAN:
```julia
# Define the hyper-parameters (here, we go with the default ones)
hparams = HyperParams()
train(hparams)
```
## Output
The generated images are stored inside the `output` folder. To visualize the output of the generator over time, we create a gif of the generated images.
```julia
folder = "output"
# Get the image filenames from the folder
img_paths = readdir(folder, join=true)
# Load all the images as an array
images = load.(img_paths)
# Join all the images in the array to create a matrix of images
gif_mat = cat(images..., dims=3)
save("./output.gif", gif_mat)
```
## Resources & References
- [The DCGAN implementation in the Model Zoo.](https://github.com/FluxML/model-zoo/blob/master/vision/dcgan_mnist/dcgan_mnist.jl)
!!! info
Originally published at [fluxml.ai](https://fluxml.ai/tutorials/) on 8 October 2021,
by Deeptendu Santra
---
### Old Tutorials/2021 10 14 Vanilla Gan
# [Tutorial: Generative Adversarial Networks](](@id man-gan-tutorial))
This tutorial describes how to implement a vanilla Generative Adversarial
Network using Flux and how train it on the MNIST dataset. It is based on this
[Pytorch tutorial](https://medium.com/ai-society/gans-from-scratch-1-a-deep-introduction-with-code-in-pytorch-and-tensorflow-cb03cdcdba0f). The original GAN [paper](https://arxiv.org/abs/1406.2661) by Goodfellow et al. is a great resource that describes the motivation and theory behind GANs:
> In the proposed adversarial nets framework, the generative model is pitted against an adversary: a
> discriminative model that learns to determine whether a sample is from the model distribution or the
> data distribution. The generative model can be thought of as analogous to a team of counterfeiters,
> trying to produce fake currency and use it without detection, while the discriminative model is
> analogous to the police, trying to detect the counterfeit currency. Competition in this game drives
> both teams to improve their methods until the counterfeits are indistinguishable from the genuine
> articles.
Let's implement a GAN in Flux. To get started we first import a few useful packages:
```julia
using MLDatasets: MNIST
using Flux.Data: DataLoader
using Flux
using CUDA
using Zygote
using UnicodePlots
```
To download a package in the Julia REPL, type `]` to enter package mode and then
type `add MLDatasets` or perform this operation with the Pkg module like this
```julia
> import Pkg
> Pkg.add("MLDatasets")
```
While [UnicodePlots](https://github.com/JuliaPlots/UnicodePlots.jl) is not necessary, it can be used to plot generated samples
into the terminal during training. Having direct feedback, instead of looking
at plots in a separate window, use fantastic for debugging.
Next, let us define values for learning rate, batch size, epochs, and other
hyper-parameters. While we are at it, we also define optimisers for the generator
and discriminator network. More on what these are later.
```julia
lr_g = 2e-4 # Learning rate of the generator network
lr_d = 2e-4 # Learning rate of the discriminator network
batch_size = 128 # batch size
num_epochs = 1000 # Number of epochs to train for
output_period = 100 # Period length for plots of generator samples
n_features = 28 * 28# Number of pixels in each sample of the MNIST dataset
latent_dim = 100 # Dimension of latent space
opt_dscr = ADAM(lr_d)# Optimiser for the discriminator
opt_gen = ADAM(lr_g) # Optimiser for the generator
```
In this tutorial I'm assuming that a CUDA-enabled GPU is available on the
system where the script is running. If this is not the case, simply remove
the `|>gpu` decorators: [piping](https://docs.julialang.org/en/v1/manual/functions/#Function-composition-and-piping).
## Data loading
The MNIST data set is available from [MLDatasets](https://juliaml.github.io/MLDatasets.jl/latest/). The first time you instantiate it you will be prompted
if you want to download it. You should agree to this.
GANs can be trained unsupervised. Therefore only keep the images from the training
set and discard the labels.
After we load the training data we re-scale the data from values in [0:1]
to values in [-1:1]. GANs are notoriously tricky to train and this re-scaling
is a recommended [GAN hack](https://github.com/soumith/ganhacks). The
re-scaled data is used to define a data loader which handles batching
and shuffling the data.
```julia
# Load the dataset
train_x, _ = MNIST.traindata(Float32);
# This dataset has pixel values ∈ [0:1]. Map these to [-1:1]
train_x = 2f0 * reshape(train_x, 28, 28, 1, :) .- 1f0 |>gpu;
# DataLoader allows to access data batch-wise and handles shuffling.
train_loader = DataLoader(train_x, batchsize=batch_size, shuffle=true);
```
## Defining the Networks
A vanilla GAN, the discriminator and the generator are both plain, [feed-forward
multilayer perceptrons](https://boostedml.com/2020/04/feedforward-neural-networks-and-multilayer-perceptrons.html). We use leaky rectified linear units [leakyrelu](https://fluxml.ai/Flux.jl/stable/models/nnlib/#NNlib.leakyrelu) to ensure out model is non-linear.
Here, the coefficient `α` (in the `leakyrelu` below), is set to 0.2. Empirically,
this value allows for good training of the network (based on prior experiments).
It has also been found that Dropout ensures a good generalization of the learned
network, so we will use that below. Dropout is usually active when training a
model and inactive in inference. Flux automatically sets the training mode when
calling the model in a gradient context. As a final non-linearity, we use the
`sigmoid` activation function.
```julia
discriminator = Chain(Dense(n_features => 1024, x -> leakyrelu(x, 0.2f0)),
Dropout(0.3),
Dense(1024 => 512, x -> leakyrelu(x, 0.2f0)),
Dropout(0.3),
Dense(512 => 256, x -> leakyrelu(x, 0.2f0)),
Dropout(0.3),
Dense(256 => 1, sigmoid)) |> gpu
```
Let's define the generator in a similar fashion. This network maps a latent
variable (a variable that is not directly observed but instead inferred) to the
image space and we set the input and output dimension accordingly. A `tanh` squashes
the output of the final layer to values in [-1:1], the same range that we squashed
the training data onto.
```julia
generator = Chain(Dense(latent_dim, 256, x -> leakyrelu(x, 0.2f0)),
Dense(256 => 512, x -> leakyrelu(x, 0.2f0)),
Dense(512 => 1024, x -> leakyrelu(x, 0.2f0)),
Dense(1024 => n_features, tanh)) |> gpu
```
## Training functions for the networks
To train the discriminator, we present it with real data from the MNIST
data set and with fake data and reward it by predicting the correct labels for
each sample. The correct labels are of course 1 for in-distribution data
and 0 for out-of-distribution data coming from the generator.
[Binary cross entropy](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.binarycrossentropy)
is the loss function of choice. While the Flux documentation suggests to use
[Logit binary cross entropy](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.logitcrossentropy),
the GAN seems to be difficult to train with this loss function.
This function returns the discriminator loss for logging purposes. We can
calculate the loss in the same call as evaluating the pullback and resort
to getting the pullback directly from Zygote instead of calling
`Flux.train!` on the model. To calculate the gradients of the loss
function with respect to the parameters of the discriminator we then only have to
evaluate the pullback with a seed gradient of 1.0. These gradients are used
to update the model parameters
```julia
function train_dscr!(discriminator, real_data, fake_data)
this_batch = size(real_data)[end] # Number of samples in the batch
# Concatenate real and fake data into one big vector
all_data = hcat(real_data, fake_data)
# Target vector for predictions: 1 for real data, 0 for fake data.
all_target = [ones(eltype(real_data), 1, this_batch) zeros(eltype(fake_data), 1, this_batch)] |> gpu;
ps = Flux.params(discriminator)
loss, pullback = Zygote.pullback(ps) do
preds = discriminator(all_data)
loss = Flux.Losses.binarycrossentropy(preds, all_target)
end
# To get the gradients we evaluate the pullback with 1.0 as a seed gradient.
grads = pullback(1f0)
# Update the parameters of the discriminator with the gradients we calculated above
Flux.update!(opt_dscr, Flux.params(discriminator), grads)
return loss
end
```
Now we need to define a function to train the generator network. The job of the
generator is to fool the discriminator so we reward the generator when the discriminator
predicts a high probability for its samples to be real data. In the training function
we first need to sample some noise, i.e. normally distributed data. This has
to be done outside the pullback since we don't want to get the gradients with
respect to the noise, but to the generator parameters. Inside the pullback we need
to first apply the generator to the noise since we will take the gradient with respect
to the parameters of the generator. We also need to call the discriminator in order
to evaluate the loss function inside the pullback. Here we need to remember to deactivate
the dropout layers of the discriminator. We do this by setting the discriminator into
test mode before the pullback. Immediately after the pullback we set it back into training
mode. Then we evaluate the pullback, call it with a seed gradient of 1.0 as above, update the
parameters of the generator network and return the loss.
```julia
function train_gen!(discriminator, generator)
# Sample noise
noise = randn(latent_dim, batch_size) |> gpu;
# Define parameters and get the pullback
ps = Flux.params(generator)
# Set discriminator into test mode to disable dropout layers
testmode!(discriminator)
# Evaluate the loss function while calculating the pullback. We get the loss for free
loss, back = Zygote.pullback(ps) do
preds = discriminator(generator(noise));
loss = Flux.Losses.binarycrossentropy(preds, 1.)
end
# Evaluate the pullback with a seed-gradient of 1.0 to get the gradients for
# the parameters of the generator
grads = back(1.0f0)
Flux.update!(opt_gen, Flux.params(generator), grads)
# Set discriminator back into automatic mode
trainmode!(discriminator, mode=:auto)
return loss
end
```
## Training
Now we are ready to train the GAN. In the training loop we keep track
of the per-sample loss of the generator and the discriminator, where
we use the batch loss returned by the two training functions defined above.
In each epoch we iterate over the mini-batches given by the data loader.
Only minimal data processing needs to be done before the training functions
can be called.
```julia
lossvec_gen = zeros(num_epochs)
lossvec_dscr = zeros(num_epochs)
for n in 1:num_epochs
loss_sum_gen = 0.0f0
loss_sum_dscr = 0.0f0
for x in train_loader
# - Flatten the images from 28x28xbatchsize to 784xbatchsize
real_data = flatten(x);
# Train the discriminator
noise = randn(latent_dim, size(x)[end]) |> gpu
fake_data = generator(noise)
loss_dscr = train_dscr!(discriminator, real_data, fake_data)
loss_sum_dscr += loss_dscr
# Train the generator
loss_gen = train_gen!(discriminator, generator)
loss_sum_gen += loss_gen
end
# Add the per-sample loss of the generator and discriminator
lossvec_gen[n] = loss_sum_gen / size(train_x)[end]
lossvec_dscr[n] = loss_sum_dscr / size(train_x)[end]
if n % output_period == 0
@show n
noise = randn(latent_dim, 4) |> gpu;
fake_data = reshape(generator(noise), 28, 4*28);
p = heatmap(fake_data, colormap=:inferno)
print(p)
end
end
```
For the hyper-parameters shown in this example, the generator produces useful
images after about 1000 epochs. And after about 5000 epochs the result look
indistinguishable from real MNIST data. Using a Nvidia V100 GPU on a 2.7
GHz Power9 CPU with 32 hardware threads, training 100 epochs takes about 80
seconds when using the GPU. The GPU utilization is between 30 and 40%.
To observe the network more frequently during training you can for example set
`output_period=20`. Training the GAN using the CPU takes about 10 minutes
per epoch and is not recommended.
## Results
Below you can see what some of the images output may look like after different numbers of epochs.
## Resources
* [A collection of GANs in Flux](https://github.com/AdarshKumar712/FluxGAN)
* [Wikipedia](https://en.wikipedia.org/wiki/Generative_adversarial_network)
* [GAN hacks](https://github.com/soumith/ganhacks)
!!! info
Originally published at [fluxml.ai](https://fluxml.ai/tutorials/) on 14 October 2021,
by Ralph Kube.
---
### Old Tutorials/2024 04 10 Blitz
# [Deep Learning with Julia & Flux: A 60 Minute Blitz](@id man-blitz)
This is a quick intro to [Flux](https://github.com/FluxML/Flux.jl) loosely based on [PyTorch's tutorial](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html). It introduces basic Julia programming, as well Zygote, a source-to-source automatic differentiation (AD) framework in Julia. We'll use these tools to build a very simple neural network.
## Arrays
The starting point for all of our models is the `Array` (sometimes referred to as a `Tensor` in other frameworks). This is really just a list of numbers, which might be arranged into a shape like a square. Let's write down an array with three elements.
```julia
x = [1, 2, 3]
```
Here's a matrix – a square array with four elements.
```julia
x = [1 2; 3 4]
```
We often work with arrays of thousands of elements, and don't usually write them down by hand. Here's how we can create an array of 5×3 = 15 elements, each a random number from zero to one.
```julia
x = rand(5, 3)
```
There's a few functions like this; try replacing `rand` with `ones`, `zeros`, or `randn` to see what they do.
By default, Julia works stores numbers is a high-precision format called `Float64`. In ML we often don't need all those digits, and can ask Julia to work with `Float32` instead. We can even ask for more digits using `BigFloat`.
```julia
x = rand(BigFloat, 5, 3)
x = rand(Float32, 5, 3)
```
We can ask the array how many elements it has.
```julia
length(x)
```
Or, more specifically, what size it has.
```julia
size(x)
```
We sometimes want to see some elements of the array on their own.
```julia
x
x[2, 3]
```
This means get the second row and the third column. We can also get every row of the third column.
```julia
x[:, 3]
```
We can add arrays, and subtract them, which adds or subtracts each element of the array.
```julia
x + x
x - x
```
Julia supports a feature called *broadcasting*, using the `.` syntax. This tiles small arrays (or single numbers) to fill bigger ones.
```julia
x .+ 1
```
We can see Julia tile the column vector `1:5` across all rows of the larger array.
```julia
zeros(5,5) .+ (1:5)
```
The `x'` syntax is used to transpose a column `1:5` into an equivalent row, and Julia will tile that across columns.
```julia
zeros(5,5) .+ (1:5)'
```
We can use this to make a times table.
```julia
(1:5) .* (1:5)'
```
Finally, and importantly for machine learning, we can conveniently do things like matrix multiply.
```julia
W = randn(5, 10)
x = rand(10)
W * x
```
Julia's arrays are very powerful, and you can learn more about what they can do [here](https://docs.julialang.org/en/v1/manual/arrays/).
### CUDA Arrays
CUDA functionality is provided separately by the [CUDA package](https://github.com/JuliaGPU/CUDA.jl). If you have a GPU and CUDA available, you can run `] add CUDA` in a REPL or IJulia to get it.
Once CUDA is loaded you can move any array to the GPU with the `cu` function, and it supports all of the above operations with the same syntax.
```julia
using CUDA
x = cu(rand(5, 3))
```
## Automatic Differentiation
You probably learned to take derivatives in school. We start with a simple mathematical function like
```julia
f(x) = 3x^2 + 2x + 1
f(5)
```
In simple cases it's pretty easy to work out the gradient by hand – here it's `6x+2`. But it's much easier to make Flux do the work for us!
```julia
using Flux: gradient
df(x) = gradient(f, x)[1]
df(5)
```
You can try this with a few different inputs to make sure it's really the same as `6x+2`. We can even do this multiple times (but the second derivative is a fairly boring `6`).
```julia
ddf(x) = gradient(df, x)[1]
ddf(5)
```
Flux's AD can handle any Julia code you throw at it, including loops, recursion and custom layers, so long as the mathematical functions you call are differentiable. For example, we can differentiate a Taylor approximation to the `sin` function.
```julia
mysin(x) = sum((-1)^k*x^(1+2k)/factorial(1+2k) for k in 0:5)
x = 0.5
mysin(x), gradient(mysin, x)
sin(x), cos(x)
```
You can see that the derivative we calculated is very close to `cos(x)`, as we expect.
This gets more interesting when we consider functions that take *arrays* as inputs, rather than just a single number. For example, here's a function that takes a matrix and two vectors (the definition itself is arbitrary)
```julia
myloss(W, b, x) = sum(W * x .+ b)
W = randn(3, 5)
b = zeros(3)
x = rand(5)
gradient(myloss, W, b, x)
```
Now we get gradients for each of the inputs `W`, `b` and `x`, which will come in handy when we want to train models.
ML models can contain hundreds of parameter arrays, therefore it is handy to group them into **layers**.
A layer is just a handy container for some parameters. For example, `Dense` does a linear transform for you.
```julia
using Flux
m = Dense(10 => 5)
x = rand(Float32, 10)
```
We can easily get the parameters of any layer or model with `trainables`.
```julia
Flux.trainables(m)
```
It very easy to calculate the gradient for all parameters in a network, even if it has many parameters.
The function `gradient` is not limited to array but can compute the gradient with respect to generic composite types.
```julia
using Flux
using Flux: logitcrossentropy, trainables, getkeypath
x = rand(Float32, 10)
model = Chain(Dense(10 => 5, relu), Dense(5 => 2))
loss(model, x) = logitcrossentropy(model(x), [0.5, 0.5])
grad = gradient(m -> loss(m, x), model)[1]
for (k, p) in trainables(model, path=true)
println("$k => $(getkeypath(grad, k))")
end
```
You don't have to use layers, but they can be convenient for many simple kinds of models and fast iteration.
The next step is to update our weights and perform optimisation. As you might be familiar, *Gradient Descent* is a simple algorithm that takes the weights and steps using a learning rate and the gradients. `weights = weights - learning_rate * gradient`.
```julia
η = 0.1
for (k, p) in trainables(model, path=true)
p .+= -η * getkeypath(grad, p)
end
```
While this is a valid way of updating our weights, it can get more complicated as the algorithms we use get more involved.
Flux comes with a bunch of pre-defined optimisers and makes writing our own really simple. We just give it the learning rate `η`:
```julia
opt_state = Flux.setup(Descent(η), model)
```
Training a network reduces down to iterating on a dataset multiple times, performing these steps in order. Just for a quick implementation, let’s train a network that learns to predict `0.5` for every input of 10 floats. `Flux` defines the `train!` function to do it for us.
```julia
data, labels = rand(10, 100), fill(0.5, 2, 100)
loss(m, x, y) = logitcrossentropy(m(x), y)
Flux.train!(loss, model, [(data, labels)], opt_state)
```
You don't have to use `train!`. In cases where arbitrary logic might be better suited, you could open up this training loop like so:
```julia
for d in training_set # assuming d looks like (data, labels)
# our super logic
g = gradient(model) do model
l = loss(model, d...)
end[1]
Flux.update!(opt_state, model, g)
end
```
The `do` block is a closure, which is a way of defining a function inline. It's a very powerful feature of Julia, and you can learn more about it [here](https://docs.julialang.org/en/v1/manual/functions/#Do-Block-Syntax-for-Function-Arguments).
## Training a Classifier
Getting a real classifier to work might help cement the workflow a bit more. [CIFAR10](https://https://www.cs.toronto.edu/~kriz/cifar.html) is a dataset of 50k tiny training images split into 10 classes.
We will do the following steps in order:
* Load CIFAR10 training and test datasets
* Define a Convolution Neural Network
* Define a loss function
* Train the network on the training data
* Test the network on the test data
### Loading the Dataset
```julia
using Statistics
using Flux
using MLDatasets: CIFAR10
using ImageCore: colorview, RGB
using Flux: onehotbatch, onecold, DataLoader
using Plots: plot
using MLUtils: splitobs, numobs
# using CUDA # Uncomment if you have CUDA installed. Can also use AMDGPU or Metal instead
# using AMDGPU
# using Metal
```
This image will give us an idea of what we are dealing with.
```julia
train_x, train_y = CIFAR10(:train)[:]
labels = onehotbatch(train_y, 0:9)
```
The `train_x` contains 50000 images converted to 32 X 32 X 3 arrays with the third dimension being the 3 channels (R,G,B). Let's take a look at a random image from the train_x. For this, we need to permute the dimensions to 3 X 32 X 32 and use `colorview` to convert it back to an image.
```julia
image(x) = colorview(RGB, permutedims(x, (3, 2, 1)))
plot(image(train_x[:,:,:,rand(1:end)]))
```
We can now arrange the training data in batches of say, 256 and keep a validation set to track our progress. This process is called minibatch learning, which is a popular method of training large neural networks. Rather that sending the entire dataset at once, we break it down into smaller chunks (called minibatches) that are typically chosen at random, and train only on them. It is shown to help with escaping [saddle points](https://en.wikipedia.org/wiki/Saddle_point).
The first 45k images (in batches of 256) will be our training set, and the rest is for validation.
The `DataLoader` function will help us load the data in batches.
```julia
trainset, valset = splitobs((train_x, labels), at = 45000)
trainloader = DataLoader(trainset, batchsize = 1000, shuffle = true)
valloader = DataLoader(trainset, batchsize = 1000)
```
### Defining the Classifier
Now we can define our Convolutional Neural Network (CNN).
A convolutional neural network is one which defines a kernel and slides it across a matrix to create an intermediate representation to extract features from. It creates higher order features as it goes into deeper layers, making it suitable for images, where the strucure of the subject is what will help us determine which class it belongs to.
```julia
model = Chain(
Conv((5,5), 3 => 16, relu),
MaxPool((2, 2)),
Conv((5, 5), 16 => 8, relu),
MaxPool((2,2)),
x -> reshape(x, :, size(x, 4)),
Dense(200 => 120),
Dense(120 => 84),
Dense(84 => 10)) |> gpu
```
We will use a crossentropy loss and an `Momentum` optimiser here. Crossentropy will be a good option when it comes to working with mulitple independent classes. Momentum gradually lowers the learning rate as we proceed with the training. It helps maintain a bit of adaptivity in our optimisation, preventing us from over shooting from our desired destination.
```julia
using Flux: logitcrossentropy, Momentum
loss(m, x, y) = logitcrossentropy(m(x), y)
opt_state = Flux.setup(Momentum(0.01), model)
```
We can start writing our train loop where we will keep track of some basic accuracy numbers about our model. We can define an `accuracy` function for it like so:
```julia
function accuracy(model, loader)
n = 0
acc = 0
for batch in loader
x, y = batch |> gpu
ŷ = model(x)
acc += sum(onecold(ŷ) .== onecold(y))
n += numobs(x)
end
return acc / n
end
```
### Training the Classifier
Training is where we do a bunch of the interesting operations we defined earlier, and see what our net is capable of. We will loop over the dataset 10 times and feed the inputs to the neural network and optimise.
```julia
epochs = 10
for epoch in 1:epochs
for batch in trainloader
x, y = batch |> gpu
g = gradient(model) do m
loss(m, x, y)
end[1]
Flux.update!(opt_state, model, g)
end
@show accuracy(model, valloader)
end
```
Seeing our training routine unfold gives us an idea of how the network learnt the function. This is not bad for a small hand-written network, trained for a limited time.
### Training on a GPU
The `gpu` functions you see sprinkled through this bit of the code tell Flux to move these entities to an available GPU, and subsequently train on it. No extra faffing about required! The same bit of code would work on any hardware with some small annotations like you saw here.
If you're not using `CUDA`, be sure to [configure Flux.jl for your GPU backend](https://fluxml.ai/Flux.jl/stable/gpu/#Selecting-GPU-backend).
### Testing the Network
We have trained the network for 100 passes over the training dataset. But we need to check if the network has learnt anything at all.
We will check this by predicting the class label that the neural network outputs, and checking it against the ground-truth. If the prediction is correct, we add the sample to the list of correct predictions. This will be done on a yet unseen section of data.
Okay, first step. Let us perform the exact same preprocessing on this set, as we did on our training set.
```julia
test_x, test_y = CIFAR10(:test)[:]
test_labels = onehotbatch(test_y, 0:9)
testloader = DataLoader((test_x, test_labels), batchsize = 1000, shuffle = true)
```
Next, display an image from the test set.
```julia
plot(image(test_x[:,:,:,rand(1:end)]))
```
The outputs of the networks are (log)likelihoods for the 10 classes. Higher the energy for a class, the more the network thinks that the image is of the particular class. Every column corresponds to the output of one image, with the 10 floats in the column being the energies.
Let's see how the model fared.
```julia
ids = rand(1:10000, 5)
rand_test = test_x[:,:,:,ids] |> gpu
rand_truth = test_y[ids]
model(rand_test)
```
This looks similar to how we would expect the results to be. At this point, it's a good idea to see how our net actually performs on new data, that we have prepared.
```julia
accuracy(model, testloader)
```
This is much better than random chance set at 10% (since we only have 10 classes), and not bad at all for a small hand written network like ours.
Let's take a look at how the net performed on all the classes performed individually.
```julia
confusion_matrix = zeros(Int, 10, 10)
m = model |> cpu
for batch in testloader
@show numobs(batch)
x, y = batch
preds = m(x)
ŷ = onecold(preds)
y = onecold(y)
for (yi, ŷi) in zip(y, ŷ)
confusion_matrix[yi, ŷi] += 1
end
end
confusion_matrix
```
The spread seems pretty good, with certain classes performing significantly better than the others. Why should that be?
!!! info
Originally published at [fluxml.ai](https://fluxml.ai/tutorials/) on 15 November 2020.
Written by Saswat Das, Mike Innes, Andrew Dinhobl, Ygor Canalli, Sudhanshu Agrawal, João Felipe Santos.
---
### Old Tutorials/2024 04 10 Mlp
# [Tutorial: Simple Multi-layer Perceptron](@id man-mlp-tutorial)
In this example, we create a simple [multi-layer perceptron](https://en.wikipedia.org/wiki/Multilayer_perceptron#:~:text=A%20multilayer%20perceptron%20(MLP)%20is,artificial%20neural%20network%20(ANN).&text=An%20MLP%20consists%20of%20at,uses%20a%20nonlinear%20activation%20function.) (MLP) that classifies handwritten digits using the MNIST dataset. A MLP consists of at least *three layers* of stacked perceptrons: Input, hidden, and output. Each neuron of an MLP has parameters (weights and bias) and uses an [activation function](https://en.wikipedia.org/wiki/Activation_function) to compute its output.
To run this example, we need the following packages:
```julia
using Flux, Statistics
using Flux: DataLoader
using Flux: onehotbatch, onecold, logitcrossentropy
# using CUDA # Uncomment this line if you have a nvidia GPU. Also AMDGPU and Metal are supported.
using MLDatasets: MNIST
using MLUtils
```
We set default values for learning rate, batch size, epochs, and the usage of a GPU (if available) for our model:
```julia
Base.@kwdef mutable struct Args
rate::Float64 = 3e-4 # learning rate
batchsize::Int = 1024 # batch size
epochs::Int = 10 # number of epochs
usegpu::Bool = true
end
```
If a GPU is available on our local system, then Flux uses it for computing the loss and updating the weights and biases when training our model.
## Data
We create the function `getdata` to load the MNIST train and test data sets from [MLDatasets](https://juliaml.github.io/MLDatasets.jl/latest/) and prepare them for the training process. In addition, we set mini-batches of the data sets by loading them onto a [DataLoader](https://fluxml.ai/Flux.jl/stable/data/dataloader/#Flux.Data.DataLoader) object.
```julia
function getdata(args)
ENV["DATADEPS_ALWAYS_ACCEPT"] = "true"
# Loading Dataset
xtrain, ytrain = MNIST(:train)[:]
xtest, ytest = MNIST(:test)[:]
# Reshape Data in order to flatten each image into a linear array
xtrain = Flux.flatten(xtrain)
xtest = Flux.flatten(xtest)
# One-hot-encode the labels
ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)
# Batching
train_loader = DataLoader((xtrain, ytrain), batchsize=args.batchsize, shuffle=true)
test_loader = DataLoader((xtest, ytest), batchsize=args.batchsize)
return train_loader, test_loader
end
```
`getdata` performs the following steps:
* **Loads MNIST data set:** Loads the train and test set tensors. The shape of train data is `28x28x60000` and test data is `28X28X10000`.
* **Reshapes the train and test data:** Uses the [flatten](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.flatten) function to reshape the train data set into a `784x60000` array and test data set into a `784x10000`. Notice that we reshape the data so that we can pass these as arguments for the input layer of our model (a simple MLP expects a vector as an input).
* **One-hot encodes the train and test labels:** Creates a batch of one-hot vectors so we can pass the labels of the data as arguments for the loss function. For this example, we use the [logitcrossentropy](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.logitcrossentropy) function and it expects data to be one-hot encoded.
* **Creates batches of data:** Creates two DataLoader objects (train and test) that handle data mini-batches of size `1024 ` (as defined above). We create these two objects so that we can pass the entire data set through the loss function at once when training our model. Also, it shuffles the data points during each iteration (`shuffle=true`).
## Model
As we mentioned above, a MLP consist of *three* layers that are fully connected. For this example, we define out model with the following layers and dimensions:
* **Input:** It has `784` perceptrons (the MNIST image size is `28x28`). We flatten the train and test data so that we can pass them as arguments to this layer.
* **Hidden:** It has `32` perceptrons that use the [relu](https://fluxml.ai/Flux.jl/stable/models/nnlib/#NNlib.relu) activation function.
* **Output:** It has `10` perceptrons that output the model's prediction or probability that a digit is 0 to 9.
We define our model with the `build_model` function:
```julia
function build_model(; imgsize=(28,28,1), nclasses=10)
return Chain(
Dense(prod(imgsize) => 32, relu),
Dense(32 => nclasses))
end
```
Note that we use the functions [Dense](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Dense) so that our model is *densely* (or fully) connected and [Chain](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Chain) to chain the computation of the three layers.
## Loss functions
Now, we define the loss function `loss_all`. It expects a DataLoader object and the `model` function we defined above as arguments. Notice that this function iterates through the `DataLoader` object in mini-batches and uses the function [logitcrossentropy](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.logitcrossentropy) to compute the difference between the predicted and actual values.
```julia
function loss_all(dataloader, model)
l = 0f0
n = 0
for (x, y) in dataloader
l += logitcrossentropy(model(x), y, agg=sum)
n += MLUtils.numobs(x)
end
return l / n
end
```
In addition, we define the function (`accuracy`) to report the accuracy of our model during the training process. To compute the accuracy, we need to decode the output of our model using the [onecold](https://fluxml.ai/Flux.jl/stable/data/onehot/#Flux.onecold) function.
```julia
function accuracy(dataloader, model)
acc = 0
n = 0
for (x, y) in dataloader
acc += sum(onecold(cpu(model(x))) .== onecold(cpu(y)))
n += MLUtils.numobs(x)
end
return acc / n
end
```
## Train our model
Finally, we create the `train` function that calls the functions we defined and trains the model.
```julia
function train(; kws...)
# Initializing Model parameters
args = Args(; kws...)
device = args.usegpu ? Flux.get_device() : Flux.get_device("CPU")
# Load Data
train_loader, test_loader = getdata(args)
# Construct model
model = build_model() |> device
loss(model, x, y) = logitcrossentropy(model(x), y)
## Training
opt_state = Flux.setup(Adam(args.rate), model)
for epoch in 1:args.epochs
@info "Epoch $epoch"
for d in train_loader
x, y = d |> device
g = gradient(m -> loss(m, x, y), model)[1]
Flux.update!(opt_state, model, g)
end
@show accuracy(train_loader, model)
@show accuracy(test_loader, model)
end
end
train()
```
`train` performs the following steps:
* **Initializes the model parameters:** Creates the `args` object that contains the default values for training our model.
* **Loads the train and test data:** Calls the function `getdata` we defined above.
* **Constructs the model:** Builds the model and loads the train and test data sets, and our model onto the GPU (if available).
* **Trains the model:** Sets [Adam](@ref Optimisers.Adam) as the optimiser for training out model, runs the training process for `10` epochs (as defined in the `args` object) and shows the `accuracy` value for the train and test data.
To see the full version of this example, see [Simple multi-layer perceptron - model-zoo](https://github.com/FluxML/model-zoo/blob/master/vision/mlp_mnist/mlp_mnist.jl).
## Resources
* [3Blue1Brown Neural networks videos](https://www.youtube.com/watch?v=aircAruvnKk&list=PLZHQObOWTQDNU6R1_67000Dx_ZCJB-3pi).
* [Neural Networks and Deep Learning](http://neuralnetworksanddeeplearning.com/).
!!! info
Originally published at [fluxml.ai](https://fluxml.ai/tutorials/) on 26 January 2021.
Written by Adarsh Kumar, Mike J Innes, Andrew Dinhobl, Jerry Ling, natema, Zhang Shitian, Liliana Badillo, Dhairya Gandhi
---
### Old Tutorials/README
These tutorials are hard to mantain
and overlapping with model-zoo examples.
Some of the tutorials are outdated.
Mantainance would be simplified by moving them
to Literate.jl and CI testing them.
---
### Src/Guide/Models/Basics
# [How Flux Works: Parameters, Gradients, and Layers](@id man-basics)
A neural network is a function with *parameters*.
That is, it takes some input `x` and gives you some output `y`,
whose value also depends on some other numbers `θ`.
A sufficiently flexible function can, by adjusting the parameters just right,
be made to do many things. And the one magic trick for adjusting parameters
is to follow a *gradient*.
This page describes Flux's take on how to construct such flexible functions
containing many parameters, and how to handle their gradients.
## Parameterised Functions
Let's start with very simple functions. This is a polynomial in `x::Real`,
returning another real number `y` which depends on some coefficients stored in a vector:
```jldoctest poly; output = false
θ = [10, 1, 0.1]
poly1(x::Real) = θ[1] + θ[2]*x + θ[3]*x^2
poly1(5) == 17.5 # true
# output
true
```
Here the parameters are a global variable `θ`. They could be handled in other ways,
for instance by explicitly passing them as an additional argument to the function:
```jldoctest poly; output = false
poly2(x::Real, θ2) = evalpoly(x, θ2) # built-in, from Base.Math
poly2(5, θ) == 17.5 # true
# output
true
```
Flux chooses a third path, by *encapsulating* the parameters within the function.
The simplest way to do this is a *closure*, an anonymous function which Julia knows
to depend on some local variable `θ3`:
```jldoctest poly; output = false
poly3 = let θ3 = [10, 1, 0.1]
x -> evalpoly(x, θ3)
end
poly3(5) == 17.5 # true
# output
true
```
An equivalent, but tidier, way is to construct a `struct` in which to store the parameters.
Any struct can be made callable, allowing its instances to act just like function:
```jldoctest poly; output = false
struct Poly3{T} # container struct
θ3::T
end
(p::Poly3)(x::Real) = evalpoly(x, p.θ3) # make this callable
poly3s = Poly3([10, 1, 0.1]) # construct an instance
poly3s(5) == 17.5 # true
# output
true
```
Internally, there is little difference between a closure and a struct.
They have the same fields, and equivalent methods:
```julia
dump(poly3), dump(poly3s) # both contain θ3: Array
poly3s.θ3 == poly3.θ3 == θ # field called :θ3 has same value
methods(poly3)
methods(poly3s) # each has 1 method, accepting x
```
The virtue of encapsulation is that it makes composition very easy.
We can make more complicated functions by combining simple ones,
and each will keep track of its own parameters.
Juia writes function composition as `∘`, for instance `(inv ∘ sin)(pi/6) ≈ 2`,
and we can use exactly this for our parameterised polynomials:
```jldoctest poly; output = false
poly4 = Poly3([1, 0.5, 0]) ∘ Poly3([10, 1, 0.1])
poly4 isa ComposedFunction # ∘ creates another struct...
poly4.outer.θ3 == θ # which has fields :inner & :outer
poly4(5) == 9.75 # true
# output
true
```
Flux models are precisely made by such function composition.
In fact, `poly3` and `poly4` are already valid Flux models.
## [Structural Gradients](@id man-taking-gradients)
The derivative of a scalar function is its slope: how fast the output changes as the input is changed slightly.
This may be found approximately by evaluating at two nearby points, and exactly by taking the limit in
which the distance between them approaches zero:
```jldoctest poly
julia> (poly1(5 + 0.1) - poly1(5)) / 0.1
2.010000000000005
julia> (poly1(5 + 0.001) - poly1(5)) / 0.001 # answer is getting close to 2
2.000100000003613
```
Flux's `gradient(f, x)` works this out for `f(x)`, and gives exactly `∂f/∂x = 2.0` here:
```jldoctest poly
julia> using Flux
julia> gradient(poly1, 5)
(2.0,)
```
The reason `gradient` returns a tuple, not just the number `2.0`, is to allow for
functions taking several arguments. (That's also why it's not called "derivative".)
For instance, this returns `∂f/∂x, ∂f/∂y, ∂f/∂z`:
```jldoctest poly
julia> gradient((x,y,z) -> (x*y)+z, 30, 40, 50)
(40.0, 30.0, 1.0)
```
For our parameterised polynomial, we have `∂f/∂x` but we are really more interested
in `∂f/∂θ`, as this will tell us about how the parameters are affecting the answer.
It is not impossible to track gradients with respect to global `θ`, but much clearer to track explicit arguments.
Here's how this works for `poly2` (which takes `θ` as a 2nd argument) and `poly3` (which encapsulates `θ`):
```jldoctest poly
julia> grad2 = gradient(poly2, 5, θ)
(2.0, [1.0, 5.0, 25.0])
julia> grad3 = gradient((x,p) -> p(x), 5, poly3s)
(2.0, (θ3 = [1.0, 5.0, 25.0],))
```
The first entry is `∂f/∂x` as before, but the second entry is more interesting.
For `poly2`, we get `∂f/∂θ` as `grad2[2]` directly.
It is a vector, because `θ` is a vector, and has elements `[∂f/∂θ[1], ∂f/∂θ[2], ∂f/∂θ[3]]`.
For `poly3s`, however, we get a `NamedTuple` whose fields correspond to those of the struct `Poly3`.
This is called a *structural gradient*. And the nice thing about them is that they work for
arbitrarily complicated structures, for instance:
```jldoctest poly
julia> grad4 = gradient(|>, 5, poly4)
(1.0, (outer = (θ3 = [1.0, 17.5, 306.25],), inner = (θ3 = [0.5, 2.5, 12.5],)))
```
Here `grad4[2].inner.θ3` corresponds to `poly4.inner.θ3`.
These matching nested structures are at the core of how Flux works.
!!! note "Implicit gradients"
Earlier versions of Flux used a different way to relate parameters and gradients,
which looks like this:
```julia
g1 = gradient(() -> poly1(5), Params([θ]))
g1[θ] == [1.0, 5.0, 25.0]
```
Here `Params` is a set of references to global variables using `objectid`,
and `g1 isa Grads` is a dictionary from these to their gradients.
This method of `gradient` takes a zero-argument function, which only *implicitly*
depends on `θ`.
## Automatic Differentiation
Flux's [`gradient`](@ref Flux.gradient) function by default calls a companion packages called [Zygote](https://github.com/FluxML/Zygote.jl).
Zygote performs source-to-source automatic differentiation, meaning that `gradient(f, x)`
hooks into Julia's compiler to find out what operations `f` contains, and transforms this
to produce code for computing `∂f/∂x`.
Zygote can in principle differentiate almost any Julia code. However, it's not perfect,
and you may eventually want to read its [page about limitations](https://fluxml.ai/Zygote.jl/dev/limitations/).
In particular, a major limitation is that mutating an array is not allowed.
Flux can also be used with other automatic differentiation (AD) packages.
It was originally written using [Tracker](https://github.com/FluxML/Tracker.jl), a more traditional operator-overloading approach.
The future might be [Enzyme](https://github.com/EnzymeAD/Enzyme.jl), and Flux now builds in an easy way to use this instead, turned on by wrapping the model in `Duplicated`. (For details, see the [Enzyme page](@ref autodiff-enzyme) in the manual.)
```julia-repl
julia> using Enzyme: Const, Duplicated
julia> grad3e = Flux.gradient((x,p) -> p(x), Const(5.0), Duplicated(poly3s))
(nothing, (θ3 = [1.0, 5.0, 25.0],))
```
`Flux.gradient` follows Zygote's convention that arguments with no derivative are marked `nothing`.
Here, this is because `Const(5.0)` is explicitly constant.
Below, we will see an example where `nothing` shows up because the model struct has fields containing things other than parameters, such as an activation function.
(It also adopts the convention that `gradient(f, x, y)` returns a tuple `(∂f/∂x, ∂f/∂y)`, without a "`∂f/∂f`" term for the function. This is why we had to write `gradient(|>, 5, poly4)` above, not just `gradient(poly4, 5)`.)
The function [`withgradient`](@ref) works the same way, but also returns the value of the function:
```jldoctest poly
julia> Flux.withgradient((x,p) -> p(x), 5.0, poly3s)
(val = 17.5, grad = (2.0, (θ3 = [1.0, 5.0, 25.0],)))
```
One can also directly specify which AD backend to use, by passing an adtype among the supported ones
(`AutoMooncake, AutoEnzyme, AutoZygote, AutoFiniteDifferences`) as the second argument.
The corresponding AD package has to be loaded first.
Here is an example using [Mooncake](https://github.com/chalk-lab/Mooncake.jl):
```jldoctest poly
julia> using Mooncake
julia> Flux.withgradient((x,p) -> p(x), AutoMooncake(), 5.0, poly3s)
(val = 17.5, grad = (2.0, (θ3 = [1.0, 5.0, 25.0],)))
```
and here is the same example using Enzyme:
```julia-repl
julia> using Enzyme
julia> Flux.withgradient((x,p) -> p(x), AutoEnzyme(), 5.0, poly3s)
(val = 17.5, grad = (2.0, Poly3{Vector{Float64}}([1.0, 5.0, 25.0])))
```
## Simple Neural Networks
The polynomial functions above send a number `x` to another a number `y`.
Neural networks typically take a vector of numbers, mix them all up, and return another vector.
Here's a very simple one, which will take a vector like `x = [1.0, 2.0, 3.0]`
and return another vector `y = layer1(x)` with `length(y) == 2`:
```jldoctest poly; output = false
W = randn(2, 3)
b = zeros(2)
sigmoid(x::Real) = 1 / (1 + exp(-x))
layer1(x) = sigmoid.(W*x .+ b)
# output
layer1 (generic function with 1 method)
```
Here `sigmoid` is a nonlinear function, applied element-wise
because it is called with `.()`, called broadcasting.
Like `poly1` above, this `layer1` has as its parameters the global variables `W, b`.
We can similarly define a version which takes these as arguments (like `poly2`),
and a version which encapsulates them (like `poly3` above):
```jldoctest poly; output = false
layer2(x, W2, b2) = sigmoid.(W2*x .+ b2) # explicit parameter arguments
layer3 = let
W3 = randn(2, 3)
b3 = zeros(2)
x -> sigmoid.(W3*x .+ b3) # closure over local variables
end
layer3([1.0, 2.0, 3.0]) isa Vector # check that it runs
# output
true
```
This third way is precisely a Flux model. And we can again make a tidier version
using a `struct` to hold the parameters:
```jldoctest poly; output = false, filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
struct Layer # container struct
W::Matrix
b::Vector
act::Function
end
(d::Layer)(x) = d.act.(d.W*x .+ d.b) # make it callabale
Layer(in::Int, out::Int, act::Function=sigmoid) =
Layer(randn(Float32, out, in), zeros(Float32, out), act)
layer3s = Layer(3, 2) # instance with its own parameters
# output
Layer(Float32[0.6911411 0.47683495 -0.75600505; 0.5247729 1.2508286 0.27635413], Float32[0.0, 0.0], sigmoid)
```
The one new thing here is a friendly constructor `Layer(in, out, act)`.
This is because we anticipate composing several instances of this thing,
with independent parameter arrays, of different sizes and different
random initial parameters.
Let's try this out, and look at its gradient:
```jldoctest poly; output = false, filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
x = Float32[0.1, 0.2, 0.3] # input
layer3s(x) # output, 2-element Vector{Float32}
Flux.gradient((x,d) -> d(x)[1], x, layer3s)[2] # NamedTuple{(:W, :b, :act)}
# output
(W = Float32[0.024975738 0.049951475 0.07492722; 0.0 0.0 0.0], b = Float32[0.24975738, 0.0], act = nothing)
```
This `∂f/∂layer3s` is a named tuple with the same fields as `Layer`.
Within it, the gradient with respect to `W` is a matrix of seemingly random numbers.
Notice that there is also an entry for `act`, which is `nothing`,
as this field of the struct is not a smoothly adjustable parameter.
We can compose these layers just as we did the polynomials above, in `poly4`.
Here's a composition of 3 functions, in which the last step is the function `only`
which takes a 1-element vector and gives us the number inside:
```jldoctest poly; output = false, filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
model1 = only ∘ Layer(20, 1, identity) ∘ Layer(1, 20)
y = model1(Float32[0.1]) # output is a Float32 number
grad = Flux.gradient(|>, [1f0], model1)[2]
# output
(outer = (outer = nothing, inner = (W = Float32[0.058179587 0.1276911 … 0.08071162 0.034993216], b = Float32[0.14223717], act = nothing)), inner = (W = Float32[-0.048111934; -0.0008379104; … ; 0.017658396; -0.015104223;;], b = Float32[-0.048111934, -0.0008379104, 0.017207285, 0.026828118, -0.024858447, -0.015956078, 0.0020494608, -0.012577536, -0.044770215, 0.01478136, 0.034534186, -0.004748393, 0.026848236, -0.016794706, -0.041044597, 0.016186379, -0.036814954, 0.034786277, 0.017658396, -0.015104223], act = nothing))
```
This gradient is starting to be a complicated nested structure.
But it works just like before: `grad.outer.inner.W` corresponds to `model1.outer.inner.W`.
We don't have to use `∘` (which makes a `ComposedFunction` struct) to combine layers.
Instead, we could define our own container struct, or use a closure.
This `model2` will work the same way (although its fields have different names):
```jldoctest poly; output = false, filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
model2 = let
lay1 = Layer(1, 20) # local variables containing layers
lay2 = Layer(20, 1, identity)
function fwd(x) # equivalent to x -> only(lay2(lay1(x)))
mid = lay1(x)
lay2(mid) |> only
end
end
model2(Float32[0.1])
Flux.gradient(|>, [1f0], model2)[2]
# output
(lay2 = (W = Float32[0.051824596 0.03971491 … 0.038365345 0.051143322], b = Float32[0.09477656], act = nothing), lay1 = (W = Float32[-0.00049770635; 0.002891017; … ; -0.0022540581; 0.0039325757;;], b = Float32[-0.00049770635, 0.002891017, -0.00865399, -0.015051818, -0.005504916, -0.004188145, -0.01533527, -0.0059600063, -0.003092169, -0.00697084, -0.012470333, -0.0048766206, -0.010671042, -0.006604657, -0.0086712, -0.0044975257, -0.0028462198, -0.009992857, -0.0022540581, 0.0039325757], act = nothing))
```
```@raw html
```
Rather than define everything from scratch every time, Flux provides a library of
commonly used layers. The same model could be defined:
```jldoctest poly; output = false
model3 = Chain(Dense(1 => 20, σ), Dense(20 => 1), only)
# output
Chain(
Dense(1 => 20, σ), # 40 parameters
Dense(20 => 1), # 21 parameters
only,
) # Total: 4 arrays, 61 parameters, 452 bytes.
```
How does this `model3` differ from the `model1` we had before?
* Flux's [`Chain`](@ref Flux.Chain) works left-to-right, the reverse of Base's `∘`.
Its contents is stored in a tuple, thus `model3.layers[1].weight` is an array.
* Flux's layer [`Dense`](@ref Flux.Dense) has only minor differences from our `struct Layer`:
- Like `struct Poly3{T}` above, it has type parameters for its fields -- the compiler does not know exactly what type `layer3s.W` will be, which costs speed.
- Its initialisation uses not `randn` (normal distribution) but [`glorot_uniform`](@ref Flux.glorot_uniform) by default.
- It reshapes some inputs (to allow several batch dimensions), and produces more friendly errors on wrong-size input.
- And it has some performance tricks: making sure element types match, and re-using some memory.
* The function [`σ`](@ref NNlib.sigmoid) is calculated in a slightly better way,
and has a rule telling Zygote how to differentiate it efficiently.
* Flux overloads `Base.show` so to give pretty printing at the REPL prompt.
Calling [`Flux.@layer Layer`](@ref Flux.@layer) will add this, and some other niceties.
All Flux layers accept a batch of samples: Instead of mapping one sample `x::Vector` to one output `y::Vector`, they map columns of a matrix `xs::Matrix` to columns of the output. This looks like `f(xs) ≈ stack(f(x) for x in eachcol(xs))` but is done more efficiently.
If what you need isn't covered by Flux's built-in layers, it's easy to write your own.
There are more details [later](@ref man-advanced), but the steps are invariably those shown for `struct Layer` above:
1. Define a `struct` which will hold the parameters.
2. Make it callable, to define how it uses them to transform the input `x`
3. Define a constructor which initialises the parameters (if the default constructor doesn't do what you want).
4. Annotate with `@layer` to opt-in to pretty printing, and other enhancements.
```@raw html
```
To deal with such nested structures, Flux relies heavily on an associated package
called Functors. Its basic function is [`fmap`](@ref Functors.fmap),
which generalises `map(f, x)` to work on almost anything.
For example, this is how [gpu](@ref Flux.gpu) moves all arrays within a model to the GPU,
reconstructing another `only ∘ Layer(...) ∘ Layer(...)` (or a `Chain` etc.) around the new `CuArray`s:
```julia
using CUDA, Functors
fmap(cu, model1)
```
And this is a very simple gradient update of the parameters, walking over `model` and `grad` simultaneously:
```julia
fmap((x, dx) -> x isa Array ? (x - dx/100) : x, model, grad)
```
!!! note
Before Flux v0.15 (and Functors v0.5), this exploration of structs was opt-in.
After defining `struct Layer` it was necessary to call `@functor Layer` (or `@layer Layer`) before Flux would look inside.
This has now changed to be opt-out: Functors (and hence Flux) will explore arbitrary structs, unless told not to (using `Functors.@leaf`).
This is why even "anonymous structs" created by closures, like `poly3` and `layer3` above, are now valid Flux models, although the use of named structs is still recommended practice.
## Curve Fitting
Above we took gradients of the output, or sometimes to the first element
of the output -- it must be a number, not a vector. Adjusting the parameters
to make this smaller won't lead us anywhere interesting. Instead, we should minimise
some *loss function* which compares the actual output to our desired output.
Perhaps the simplest example is curve fitting. The [previous page](@ref man-overview) fitted
a linear model to data. With our two-layer model, we can fit a nonlinear function.
For example, let us use `f(x) = 2x - x^3` evaluated at some points `x in -2:0.1:2` as the data,
and adjust the parameters of `model3` from above so that its output is similar.
```jldoctest poly; output = false
data = [([x], 2x-x^3) for x in -2:0.1f0:2] # training points (x, y)
for _ in 1:1000 # adjust parameters to minimise the error:
Flux.train!((m,x,y) -> (m(x) - y)^2, model3, data, Descent(0.01))
end
# output
```
The same code will also work with `model1` or `model2` instead.
Here's how to plot the desired and actual outputs:
```julia
using Plots
plot(x -> 2x-x^3, -2, 2, label="truth")
scatter!(x -> model3([x]), -2:0.1f0:2, label="fitted")
```
More detail about what exactly the function `train!` is doing, and how to use rules other than simple [`Descent`](@ref Optimisers.Descent), is what the next page in this guide is about: [training](@ref man-training).
---
### Src/Guide/Models/Overview
# [Flux Overview: Fitting a Straight Line](@id man-overview)
Flux is a pure Julia ML stack that allows you to build predictive models. Here are the steps for a typical Flux program:
1. Provide training and test data
2. Build a model with configurable *parameters* to make predictions
3. Iteratively train the model by tweaking the parameters to improve predictions
4. Verify your model
Under the hood, Flux uses a technique called automatic differentiation to take gradients that help improve predictions. Flux is also fully written in Julia so you can easily replace any layer of Flux with your own code to improve your understanding or satisfy special requirements.
Here's how you'd use Flux to build and train the most basic of models, step by step.
### A Trivial Prediction
This example will predict the output of the function `4x + 2`. Making such predictions is called "linear regression", and is really too simple to *need* a neural network. But it's a nice toy example.
First, import `Flux` and define the function we want to simulate:
```jldoctest overview
julia> using Flux
julia> actual(x) = 4x + 2
actual (generic function with 1 method)
```
This example will build a model to approximate the `actual` function.
## 1. Provide Training and Test Data
Use the `actual` function to build sets of data for training and verification:
```jldoctest overview
julia> x_train, x_test = hcat(0:5...), hcat(6:10...)
([0 1 … 4 5], [6 7 … 9 10])
julia> y_train, y_test = actual.(x_train), actual.(x_test)
([2 6 … 18 22], [26 30 … 38 42])
```
Normally, your training and test data come from real world observations, but here we simulate them.
## 2. Build a Model to Make Predictions
Now, build a model to make predictions with `1` input and `1` output:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> model = Dense(1 => 1)
Dense(1 => 1) # 2 parameters
julia> model.weight
1×1 Matrix{Float32}:
0.95041317
julia> model.bias
1-element Vector{Float32}:
0.0
```
Under the hood, a dense layer is a struct with fields `weight` and `bias`. `weight` represents a weights' matrix and `bias` represents a bias vector. There's another way to think about a model. In Flux, *models are conceptually predictive functions*:
```jldoctest overview
julia> predict = Dense(1 => 1)
Dense(1 => 1) # 2 parameters
```
`Dense(1 => 1)` also implements the function `σ(Wx+b)` where `W` and `b` are the weights and biases. `σ` is an activation function (more on activations later). Our model has one weight and one bias, but typical models will have many more. Think of weights and biases as knobs and levers Flux can use to tune predictions. Activation functions are transformations that tailor models to your needs.
This model will already make predictions, though not accurate ones yet:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> predict(x_train)
1×6 Matrix{Float32}:
0.0 0.906654 1.81331 2.71996 3.62662 4.53327
```
In order to make better predictions, you'll need to provide a *loss function* to tell Flux how to objectively *evaluate* the quality of a prediction. Loss functions compute the cumulative distance between actual values and predictions.
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> using Statistics
julia> loss(model, x, y) = mean(abs2.(model(x) .- y));
julia> loss(predict, x_train, y_train)
122.64734f0
```
More accurate predictions will yield a lower loss. You can write your own loss functions or rely on those already provided by Flux. This loss function is called [mean squared error](https://www.statisticshowto.com/probability-and-statistics/statistics-definitions/mean-squared-error/) (and built-in as [`mse`](@ref Flux.Losses.mse)). Flux works by iteratively reducing the loss through *training*.
## 3. Improve the Prediction
Under the hood, the Flux [`Flux.train!`](@ref) function uses *a loss function* and *training data* to improve the *parameters* of your model based on a pluggable [`optimiser`](../../reference/training/optimisers.md):
```jldoctest overview
julia> using Flux: train!
julia> opt = Descent()
Descent(0.1f0)
julia> data = [(x_train, y_train)]
1-element Vector{Tuple{Matrix{Int64}, Matrix{Int64}}}:
([0 1 … 4 5], [2 6 … 18 22])
```
Now, we have the optimiser and data we'll pass to `train!`. All that remains are the parameters of the model. Remember, each model is a Julia struct with a function and configurable parameters. Remember, the dense layer has weights and biases that depend on the dimensions of the inputs and outputs:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> predict.weight
1×1 Matrix{Float32}:
0.9066542
julia> predict.bias
1-element Vector{Float32}:
0.0
```
The dimensions of these model parameters depend on the number of inputs and outputs.
Flux will adjust predictions by iteratively changing these parameters according to the optimiser.
This optimiser implements the classic gradient descent strategy. Now improve the parameters of the model with a call to [`Flux.train!`](@ref) like this:
```jldoctest overview
julia> train!(loss, predict, data, opt)
```
And check the loss:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> loss(predict, x_train, y_train)
116.38745f0
```
It went down. Why?
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> predict.weight, predict.bias
(Float32[7.246838;;], Float32[1.748103])
```
The parameters have changed. This single step is the essence of machine learning.
## 3+. Iteratively Train the Model
In the previous section, we made a single call to `train!` which iterates over the data we passed in just once. An *epoch* refers to one pass over the dataset. Typically, we will run the training for multiple epochs to drive the loss down even further. Let's run it a few more times:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> for epoch in 1:200
train!(loss, predict, data, opt)
end
julia> loss(predict, x_train, y_train)
0.00339581f0
julia> predict.weight, predict.bias
(Float32[4.0159144;;], Float32[2.004479])
```
After 200 training steps, the loss went down, and the parameters are getting close to those in the function the model is built to predict.
## 4. Verify the Results
Now, let's verify the predictions:
```jldoctest overview; filter = r"[+-]?([0-9]*[.])?[0-9]+(f[+-]*[0-9])?"
julia> predict(x_test)
1×5 Matrix{Float32}:
26.1121 30.13 34.1479 38.1657 42.1836
julia> y_test
1×5 Matrix{Int64}:
26 30 34 38 42
```
The predictions are good. Here's how we got there.
First, we gathered real-world data into the variables `x_train`, `y_train`, `x_test`, and `y_test`. The `x_*` data defines inputs, and the `y_*` data defines outputs. The `*_train` data is for training the model, and the `*_test` data is for verifying the model. Our data was based on the function `4x + 2`.
Then, we built a single input, single output predictive model, `predict = Dense(1 => 1)`. The initial predictions weren't accurate, because we had not trained the model yet.
After building the model, we trained it with `train!(loss, predict, data, opt)`. The loss function is first, followed by the model itself, the training data, and the `Descent` optimiser provided by Flux. We ran the training step once, and observed that the parameters changed and the loss went down. Then, we ran the `train!` many times to finish the training process.
After we trained the model, we verified it with the test data to verify the results.
This overall flow represents how Flux works. Let's drill down a bit to understand what's going on inside the individual layers of Flux.
---
### Src/Guide/Models/Quickstart
# [A Neural Network in One Minute](@id man-quickstart)
If you have used neural networks before, then this simple example might be helpful for seeing how the major parts of Flux work together. Try pasting the code into the REPL prompt.
If you haven't, then you might prefer the [Fitting a Straight Line](overview.md) page.
```
/* Detailed source-code truncated for AI context efficiency. */
```
```julia
using Plots # to draw the above figure
p_true = scatter(noisy[1,:], noisy[2,:], zcolor=truth, title="True classification", legend=false)
p_raw = scatter(noisy[1,:], noisy[2,:], zcolor=probs1[1,:], title="Untrained network", label="", clims=(0,1))
p_done = scatter(noisy[1,:], noisy[2,:], zcolor=probs2[1,:], title="Trained network", legend=false)
plot(p_true, p_raw, p_done, layout=(1,3), size=(1000,330))
```
```@raw html
```
Here's the loss during training:
```julia
plot(losses; xaxis=(:log10, "iteration"),
yaxis="loss", label="per batch")
n = length(loader)
plot!(n:n:length(losses), mean.(Iterators.partition(losses, n)),
label="epoch mean", dpi=200)
```
This XOR ("exclusive or") problem is a variant of the famous one which drove Minsky and Papert to invent deep neural networks in 1969. For small values of "deep" -- this has one hidden layer, while earlier perceptrons had none. (What they call a hidden layer, Flux calls the output of the first layer, `model[1](noisy)`.)
Since then things have developed a little.
## Features to Note
Some things to notice in this example are:
* The batch dimension of data is always the last one. Thus a `2×1000 Matrix` is a thousand observations, each a column of length 2. Flux defaults to `Float32`, but most of Julia to `Float64`.
* The `model` can be called like a function, `y = model(x)`. Each layer like [`Dense`](@ref Flux.Dense) is an ordinary `struct`, which encapsulates some arrays of parameters (and possibly other state, as for [`BatchNorm`](@ref Flux.BatchNorm)).
* But the model does not contain the loss function, nor the optimisation rule. The momenta needed by [`Adam`](@ref Optimisers.Adam) are stored in the object returned by [setup](@ref Flux.Train.setup). And [`Flux.logitcrossentropy`](@ref Flux.Losses.logitcrossentropy) is an ordinary function that combines the [`softmax`](@ref Flux.softmax) and [`crossentropy`](@ref Flux.crossentropy) functions.
* The `do` block creates an anonymous function, as the first argument of `gradient`. Anything executed within this is differentiated.
Instead of calling [`gradient`](@ref Flux.gradient) and [`update!`](@ref Flux.update!) separately, there is a convenience function [`train!`](@ref Flux.train!). If we didn't want anything extra (like logging the loss), we could replace the training loop with the following:
```julia
for epoch in 1:1_000
Flux.train!(model, loader |> device, opt_state) do m, x, y
y_hat = m(x)
Flux.logitcrossentropy(y_hat, y)
end
end
```
If you want to keep control of the loop — for logging, early stopping, or a custom schedule — but would rather not spell out `gradient` and `update!` yourself, [`trainstep!`](@ref Flux.Train.trainstep!) is the per-step primitive that `train!` is built on. It evaluates the loss, computes the gradient, and updates `model` and `opt_state` in place, returning the loss. The batch is passed as a tuple and spliced into the loss after the model (so it is called as `loss(m, x, y)`):
```julia
losses = []
@showprogress for epoch in 1:1_000
for xy_cpu in loader
x, y = xy_cpu |> device
loss = Flux.trainstep!(model, (x, y), opt_state) do m, x, y
y_hat = m(x)
Flux.logitcrossentropy(y_hat, y)
end
push!(losses, loss)
end
end
```
This is exactly the original training loop above, with the `withgradient`/`update!` pair collapsed into a single call. Use [`trainstep_withgradient!`](@ref Flux.Train.trainstep_withgradient!) if you also need the gradient. On a [Reactant](../reactant.md) device, `trainstep!` (and `train!`) additionally compile and cache the whole step into one executable.
* Notice that the full dataset `noisy` lives on the CPU, and is moved to the GPU one batch at a time, by `xy_cpu |> device`. This is generally what you want for large datasets. Calling `loader |> device` similarly modifies the `DataLoader` to move one batch at a time.
* In our simple example, we conveniently created the model has a [`Chain`](@ref Flux.Chain) of layers.
For more complex models, you can define a custom struct `MyModel` containing layers and arrays and implement the call operator `(::MyModel)(x) = ...` to define the forward pass. This is all that is needed for Flux to work. Marking the struct with [`Flux.@layer`](@ref) will add some more functionality, like pretty printing and the ability to mark some internal fields as trainable or not (also see [`trainable`](@ref Optimisers.trainable)).
---
### Src/Guide/Models/Recurrence
# Recurrent Models
## Recurrent cells
To introduce Flux's recurrence functionalities, we will consider the following vanilla recurrent neural network structure:
In the above, we have a sequence of length 3, where `x1` to `x3` represent the input at each step. It could be a timestamp or a word in a sentence encoded as vectors. `y1` to `y3` are their respective outputs.
An aspect to recognise is that in such a model, the recurrent cells `A` all refer to the same structure. What distinguishes it from a simple dense layer is that the cell `A` is fed, in addition to an input `x`, with information from the previous state of the model (hidden state denoted as `h1` & `h2` in the diagram).
In the most basic RNN case, cell A could be defined by the following:
```julia
output_size = 5
input_size = 2
Wxh = randn(Float32, output_size, input_size)
Whh = randn(Float32, output_size, output_size)
b = zeros(Float32, output_size)
function rnn_cell(x, h)
h = tanh.(Wxh * x .+ Whh * h .+ b)
return h, h
end
seq_len = 3
# dummy input data
x = [rand(Float32, input_size) for i = 1:seq_len]
# random initial hidden state
h0 = zeros(Float32, output_size)
y = []
ht = h0
for xt in x
yt, ht = rnn_cell(xt, ht)
y = [y; [yt]] # concatenate in non-mutating (AD friendly) way
end
```
Notice how the above is essentially a `Dense` layer that acts on two inputs, `xt` and `ht`.
The result of the forward pass at each time step, is a tuple containing the output `yt` and the updated state `ht`. The updated state is used as an input in next iteration. In the simple case of a vanilla RNN, the
output and the state are the same. In more complex cells, such as `LSTMCell`, the state can contain multiple arrays.
There are various recurrent cells available in Flux, notably `RNNCell`, `LSTMCell` and `GRUCell`, which are documented in the [layer reference](../../reference/models/layers.md). The hand-written example above can be replaced with:
```julia
using Flux
output_size = 5
input_size = 2
seq_len = 3
x = [rand(Float32, input_size) for i = 1:seq_len]
h0 = zeros(Float32, output_size)
rnn_cell = Flux.RNNCell(input_size => output_size)
y = []
ht = h0
for xt in x
yt, ht = rnn_cell(xt, ht)
y = [y; [yt]]
end
```
The entire output `y` or just the last output `y[end]` can be used for further processing, such as classification or regression.
## Using a cell as part of a model
Let's consider a simple model that is trained to predict a scalar quantity for each time step in a sequence. The model will have a single RNN cell, followed by a dense layer to produce the output.
Since the [`RNNCell`](@ref) can deal with batches of data, we can define the model to accept an input where
at each time step, the input is a matrix of size `(input_size, batch_size)`.
```julia
struct RecurrentCellModel{H,C,D}
h0::H
cell::C
dense::D
end
# we choose to not train the initial hidden state
Flux.@layer RecurrentCellModel trainable=(cell, dense)
function RecurrentCellModel(input_size::Int, hidden_size::Int)
return RecurrentCellModel(
zeros(Float32, hidden_size),
RNNCell(input_size => hidden_size),
Dense(hidden_size => 1))
end
function (m::RecurrentCellModel)(x)
z = []
ht = m.h0
for xt in x
yt, ht = m.cell(xt, ht)
z = [z; [yt]]
end
z = stack(z, dims=2) # [hidden_size, seq_len, batch_size] or [hidden_size, seq_len]
ŷ = m.dense(z) # [1, seq_len, batch_size] or [1, seq_len]
return ŷ
end
```
Notice that we stack the hidden states `z` to form a tensor of size `(hidden_size, seq_len, batch_size)`. This can speed up the final classification, since we then process all the outputs at once with a single forward pass of the dense layer.
Let's now define the training loop for this model:
```julia
using Optimisers: AdamW
function loss(model, x, y)
ŷ = model(x)
y = stack(y, dims=2)
return Flux.mse(ŷ, y)
end
# create dummy data
seq_len, batch_size, input_size = 3, 4, 2
x = [rand(Float32, input_size, batch_size) for _ = 1:seq_len]
y = [rand(Float32, 1, batch_size) for _ = 1:seq_len]
# initialize the model and optimizer
model = RecurrentCellModel(input_size, 5)
opt_state = Flux.setup(AdamW(1e-3), model)
# compute the gradient and update the model
g = gradient(m -> loss(m, x, y), model)[1]
Flux.update!(opt_state, model, g)
```
## Handling the whole sequence at once
In the above example, we processed the sequence one time step at a time using a recurrent cell. However, it is possible to process the entire sequence at once. This can be done by stacking the input data `x` to form a tensor of size `(input_size, seq_len)` or `(input_size, seq_len, batch_size)`.
One can then use the [`RNN`](@ref), [`LSTM`](@ref) or [`GRU`](@ref) layers to process the entire input tensor.
Let's consider the same example as above, but this time we use an `RNN` layer instead of an `RNNCell`:
```julia
struct RecurrentModel{H,C,D}
h0::H
rnn::C
dense::D
end
Flux.@layer RecurrentModel trainable=(rnn, dense)
function RecurrentModel(input_size::Int, hidden_size::Int)
return RecurrentModel(
zeros(Float32, hidden_size),
RNN(input_size => hidden_size),
Dense(hidden_size => 1))
end
function (m::RecurrentModel)(x)
z = m.rnn(x, m.h0) # [hidden_size, seq_len, batch_size] or [hidden_size, seq_len]
ŷ = m.dense(z) # [1, seq_len, batch_size] or [1, seq_len]
return ŷ
end
seq_len, batch_size, input_size = 3, 4, 2
x = rand(Float32, input_size, seq_len, batch_size)
y = rand(Float32, 1, seq_len, batch_size)
model = RecurrentModel(input_size, 5)
opt_state = Flux.setup(AdamW(1e-3), model)
g = gradient(m -> Flux.mse(m(x), y), model)[1]
Flux.update!(opt_state, model, g)
```
Finally, the [`Recurrence`](@ref) layer can be used wrap any recurrent cell to process the entire sequence at once. For instance, a type behaving the same as the `LSTM` layer can be defined as follows:
```julia
rnn = Recurrence(LSTMCell(2 => 3)) # similar to LSTM(2 => 3)
x = rand(Float32, 2, 4, 3)
y = rnn(x)
```
## Stacking recurrent layers
Recurrent layers can be stacked to form a deeper model by simply chaining them together using the [`Chain`](@ref) layer. The output of a layer is fed as input to the next layer in the chain.
For instance, a model with two LSTM layers can be defined as follows:
```julia
stacked_rnn = Chain(LSTM(3 => 5), Dropout(0.5), LSTM(5 => 5))
x = rand(Float32, 3, 4)
y = stacked_rnn(x)
```
If more fine grained control is needed, for instance to have a trainable initial hidden state, one can define a custom model as follows:
```julia
struct StackedRNN{L,S}
layers::L
states0::S
end
Flux.@layer StackedRNN
function StackedRNN(d::Int; num_layers::Int)
layers = [LSTM(d => d) for _ in 1:num_layers]
states0 = [Flux.initialstates(l) for l in layers]
return StackedRNN(layers, states0)
end
function (rnn::StackedRNN)(x)
for (layer, state0) in zip(rnn.layers, rnn.states0)
x = layer(x, state0)
end
return x
end
rnn = StackedRNN(3; num_layers=2)
x = rand(Float32, 3, 10)
y = rnn(x)
```
---
### Src/Guide/Training/Training
# [Training a Flux Model](@id man-training)
Training refers to the process of slowly adjusting the parameters of a model to make it work better.
Besides the model itself, we will need three things:
* An *objective function* that evaluates how well a model is doing on some input.
* An *optimisation rule* which describes how the model's parameters should be adjusted.
* Some *training data* to use as the input during this process.
Usually the training data is some collection of examples (or batches of examples) which
are handled one-by-one. One *epoch* of training means that each example is used once,
something like this:
```julia
# Initialise the optimiser for this model:
opt_state = Flux.setup(rule, model)
for data in train_set
# Unpack this element (for supervised training):
input, label = data
# Calculate the gradient of the objective
# with respect to the parameters within the model:
grads = Flux.gradient(model) do m
result = m(input)
loss(result, label)
end
# Update the parameters so as to reduce the objective,
# according the chosen optimisation rule:
Flux.update!(opt_state, model, grads[1])
end
```
This loop can also be written using the function [`train!`](@ref Flux.Train.train!),
but it's helpful to understand the pieces first:
```julia
train!(model, train_set, opt_state) do m, x, y
loss(m(x), y)
end
```
## Model Gradients
Fist recall from the section on [taking gradients](@ref man-taking-gradients) that
`Flux.gradient(f, a, b)` always calls `f(a, b)`, and returns a tuple `(∂f_∂a, ∂f_∂b)`.
In the code above, the function `f` passed to `gradient` is an anonymous function with
one argument, created by the `do` block, hence `grads` is a tuple with one element.
Instead of a `do` block, we could have written:
```julia
grads = Flux.gradient(m -> loss(m(input), label), model)
```
Since the model is some nested set of layers, `grads[1]` is a similarly nested set of
`NamedTuple`s, ultimately containing gradient components. If (for example)
`θ = model.layers[1].weight[2,3]` is one scalar parameter, an entry in a matrix of weights,
then the derivative of the loss with respect to it is `∂f_∂θ = grads[1].layers[1].weight[2,3]`.
It is important that the execution of the model takes place inside the call to `gradient`,
in order for the influence of the model's parameters to be observed by Zygote.
It is also important that every `update!` step receives a newly computed gradient,
as it will change whenever the model's parameters are changed, and for each new data point.
## Loss Functions
The objective function must return a number representing how far the model is from
the desired result. This is termed the *loss* of the model.
This number can be produced by any ordinary Julia code, but this must be executed
within the call to `gradient`. For instance, we could define a function
```julia
loss(y_hat, y) = sum((y_hat .- y).^2)
```
or write this directly inside the `do` block above. Many commonly used functions,
like [`mse`](@ref Flux.Losses.mse) for mean-squared error or [`crossentropy`](@ref Flux.Losses.crossentropy) for cross-entropy loss,
are available from the [`Flux.Losses`](../../reference/models/losses.md) module.
## Optimisation Rules
The simplest kind of optimisation using the gradient is termed *gradient descent*
(or sometimes *stochastic gradient descent* when, as here, it is not applied to the entire dataset at once).
Gradient descent needs a *learning rate* which is a small number describing how fast to walk downhill,
usually written as the Greek letter "eta", `η`. This is often described as a *hyperparameter*,
to distinguish it from the parameters which are being updated `θ = θ - η * ∂loss_∂θ`.
We want to update all the parameters in the model, like this:
```julia
η = 0.01 # learning rate
# For each parameter array, update
# according to the corresponding gradient:
fmap(model, grads[1]) do p, g
p .= p .- η .* g
end
```
A slightly more refined version of this loop to update all the parameters is wrapped up as a function [`update!`](@ref Optimisers.update!)`(opt_state, model, grads[1])`.
And the learning rate is the only thing stored in the [`Descent`](@ref Optimisers.Descent) struct.
However, there are many other optimisation rules, which adjust the step size and
direction in various clever ways.
Most require some memory of the gradients from earlier steps, rather than always
walking straight downhill -- [`Momentum`](@ref Optimisers.Momentum) is the simplest.
The function [`setup`](@ref Flux.Train.setup) creates the necessary storage for this, for a particular model.
It should be called once, before training, and returns a tree-like object which is the
first argument of `update!`. Like this:
```julia
# Initialise momentum
opt_state = Flux.setup(Momentum(0.01, 0.9), model)
for data in train_set
grads = [...]
# Update both model parameters and optimiser state:
Flux.update!(opt_state, model, grads[1])
end
```
Many commonly-used optimisation rules, such as [`Adam`](@ref Optimisers.Adam), are built-in.
These are listed on the [optimisers](@ref man-optimisers) page.
!!! compat "Implicit-style optimiser state"
This `setup` makes another tree-like structure. Old versions of Flux did not do this,
and instead stored a dictionary-like structure within the optimiser `Adam(0.001)`.
This was initialised on first use of the version of `update!` for "implicit" parameters.
## Datasets & Batches
The loop above iterates through `train_set`, expecting at each step a tuple `(input, label)`.
The very simplest such object is a vector of tuples, such as this:
```julia
x = randn(28, 28)
y = rand(10)
data = [(x, y)]
```
or `data = [(x, y), (x, y), (x, y)]` for the same values three times.
Very often, the initial data is large arrays which you need to slice into examples.
To produce one iterator of pairs `(x, y)`, you might want `zip`:
```julia
X = rand(28, 28, 60_000); # many images, each 28 × 28
Y = rand(10, 60_000)
data = zip(eachslice(X; dims=3), eachcol(Y))
first(data) isa Tuple{AbstractMatrix, AbstractVector} # true
```
Here each iteration will use one matrix `x` (an image, perhaps) and one vector `y`.
It is very common to instead train on *batches* of such inputs (or *mini-batches*,
the two words mean the same thing) both for efficiency and for better results.
This can be easily done using the [`DataLoader`](@ref Flux.DataLoader):
```julia
data = Flux.DataLoader((X, Y), batchsize=32)
x1, y1 = first(data)
size(x1) == (28, 28, 32)
length(data) == 1875 === 60_000 ÷ 32
```
Flux's layers are set up to accept such a batch of input data,
and the convolutional layers such as [`Conv`](@ref Flux.Conv) require it.
The batch index is always the last dimension.
## Training Loops
Simple training loops like the one above can be written compactly using
the [`train!`](@ref Flux.Train.train!) function. Including `setup`, this reads:
```julia
opt_state = Flux.setup(Adam(), model)
for epoch in 1:100
Flux.train!(model, train_set, opt_state) do m, x, y
loss(m(x), y)
end
end
```
Or explicitly writing the anonymous function which this `do` block creates,
`train!((m,x,y) -> loss(m(x),y), model, train_set, opt_state)` is exactly equivalent.
If you want to write the loop yourself but keep the convenience of a single gradient +
update call, [`trainstep!`](@ref Flux.Train.trainstep!) is the per-step primitive that
`train!` is built on. It differentiates the loss, updates `model` and `opt_state` in place,
and returns the loss:
```julia
opt_state = Flux.setup(Adam(), model)
for epoch in 1:100
for (x, y) in train_set
l = Flux.trainstep!((m, x, y) -> loss(m(x), y), model, (x, y), opt_state)
end
end
```
Use [`trainstep_withgradient!`](@ref Flux.Train.trainstep_withgradient!) if you also need the
gradient (it returns `(loss, grad)`). On a Reactant device both compile and cache the whole step,
just like `train!`.
Like [`withgradient`](@ref Flux.withgradient), the loss passed to `trainstep!` may return auxiliary
data alongside the scalar loss — return a `Tuple` or `NamedTuple` whose first element is the loss,
and the gradient is taken of the loss alone while the whole value is returned. This is convenient
for logging a metric computed during the forward pass:
```julia
l, stats = Flux.trainstep!(model, (x, y), opt_state) do m, x, y
ŷ = m(x)
loss(ŷ, y), (; acc = accuracy(ŷ, y))
end
```
Real training loops often need more flexibility, and the best way to do this is just
to write the loop. This is ordinary Julia code, without any need to work through some
callback API. Here is an example, in which it may be helpful to note:
* The function [`withgradient`](@ref Zygote.withgradient) is like `gradient` but also
returns the value of the function, for logging or diagnostic use.
* Logging or printing is best done outside of the `gradient` call,
as there is no need to differentiate these commands.
* To use `result` for logging purposes, you could change the `do` block to end with
`return my_loss(result, label), result`, i.e. make the function passed to `withgradient`
return a tuple. The first element is always the loss.
* Julia's `break` and `continue` keywords let you exit from parts of the loop.
```julia
opt_state = Flux.setup(Adam(), model)
my_log = []
for epoch in 1:100
losses = Float32[]
for (i, data) in enumerate(train_set)
input, label = data
val, grads = Flux.withgradient(model) do m
# Any code inside here is differentiated.
# Evaluation of the model and loss must be inside!
result = m(input)
my_loss(result, label)
end
# Save the loss from the forward pass. (Done outside of gradient.)
push!(losses, val)
# Detect loss of Inf or NaN. Print a warning, and then skip update!
if !isfinite(val)
@warn "loss is $val on item $i" epoch
continue
end
Flux.update!(opt_state, model, grads[1])
end
# Compute some accuracy, and save details as a NamedTuple
acc = my_accuracy(model, train_set)
push!(my_log, (; acc, losses))
# Stop training when some criterion is reached
if acc > 0.95
println("stopping after $epoch epochs")
break
end
end
```
## Regularisation
The term *regularisation* covers a wide variety of techniques aiming to improve the
result of training. This is often done to avoid overfitting.
Some of these can be implemented by simply modifying the loss function.
*L₂ regularisation* (sometimes called ridge regression) adds to the loss a penalty
proportional to `θ^2` for every scalar parameter.
A very simple model could be implemented as follows:
```julia
grads = Flux.gradient(densemodel) do m
result = m(input)
penalty = sum(abs2, m.weight)/2 + sum(abs2, m.bias)/2
my_loss(result, label) + 0.42f0 * penalty
end
```
Accessing each individual parameter array by hand won't work well for large models.
Instead, we can use [`Flux.trainables`](@ref Optimisers.trainables) to collect all of them,
and then apply a function to each one, and sum the result:
```julia
pen_l2(x::AbstractArray) = sum(abs2, x)/2
grads = Flux.gradient(model) do m
result = m(input)
penalty = sum(pen_l2, Flux.trainables(m))
my_loss(result, label) + 0.42f0 * penalty
end
```
However, the gradient of this penalty term is very simple: It is proportional to the original weights.
So there is a simpler way to implement exactly the same thing, by modifying the optimiser
instead of the loss function. This is done by replacing this:
```julia
opt_state = Flux.setup(Adam(0.1), model)
```
with this:
```julia
decay_opt_state = Flux.setup(OptimiserChain(WeightDecay(0.42), Adam(0.1)), model)
```
Flux's optimisers are really modifications applied to the gradient before using it to update
the parameters, and [`OptimiserChain`](@ref Optimisers.OptimiserChain) applies two such modifications.
The first, [`WeightDecay`](@ref Optimisers.WeightDecay) adds `0.42` times the original parameter to the gradient,
matching the gradient of the penalty above (with the same, unrealistically large, constant).
After that, in either case, [`Adam`](@ref Optimisers.Adam) computes the final update.
The same trick works for *L₁ regularisation* (also called Lasso), where the penalty is
`pen_l1(x::AbstractArray) = sum(abs, x)` instead. This is implemented by `SignDecay(0.42)`.
The same `OptimiserChain` mechanism can be used for other purposes, such as gradient clipping with [`ClipGrad`](@ref Optimisers.ClipGrad) or [`ClipNorm`](@ref Optimisers.ClipNorm).
Besides L1 / L2 / weight decay, another common and quite different kind of regularisation is
provided by the [`Dropout`](@ref Flux.Dropout) layer. This turns off some outputs of the
previous layer during training.
It should switch automatically, but see [`trainmode!`](@ref Flux.trainmode!) / [`testmode!`](@ref Flux.testmode!) to manually enable or disable this layer.
## Learning Rate Schedules
Finer control of training, you may wish to alter the learning rate mid-way through training.
This can be done with [`adjust!`](@ref Flux.adjust!), like this:
```julia
opt_state = Flux.setup(Adam(0.1), model) # initialise once
for epoch in 1:1000
train!([...], state) # Train with η = 0.1 for first 100,
if epoch == 100 # then change to use η = 0.01 for the rest.
Flux.adjust!(opt_state, 0.01)
end
end
```
Other hyper-parameters can also be adjusted, such as `Flux.adjust!(opt_state, beta = (0.8, 0.99))`.
And such modifications can be applied to just one part of the model.
For instance, this sets a different learning rate for the encoder and the decoder:
```julia
# Consider some model with two parts:
bimodel = Chain(enc = [...], dec = [...])
# This returns a tree whose structure matches the model:
opt_state = Flux.setup(Adam(0.02), bimodel)
# Adjust the learning rate to be used for bimodel.layers.enc
Flux.adjust!(opt_state.layers.enc, 0.03)
```
## Scheduling Optimisers
In practice, it is fairly common to schedule the learning rate of an optimiser to obtain faster convergence. There are a variety of popular scheduling policies, and you can find implementations of them in [ParameterSchedulers.jl](http://fluxml.ai/ParameterSchedulers.jl/stable). The documentation for ParameterSchedulers.jl provides a more detailed overview of the different scheduling policies, and how to use them with Flux optimisers. Below, we provide a brief snippet illustrating a [cosine annealing](https://arxiv.org/pdf/1608.03983.pdf) schedule with a momentum optimiser.
First, we import ParameterSchedulers.jl and initialize a cosine annealing schedule to vary the learning rate between `1e-4` and `1e-2` every 10 epochs. We also create a new [`Momentum`](@ref Optimisers.Momentum) optimiser.
```julia
using ParameterSchedulers
opt_state = Flux.setup(Momentum(), model)
schedule = Cos(λ0 = 1e-4, λ1 = 1e-2, period = 10)
for (eta, epoch) in zip(schedule, 1:100)
Flux.adjust!(opt_state, eta)
# your training code here
end
```
`schedule` can also be indexed (e.g. `schedule(100)`) or iterated like any iterator in Julia.
ParameterSchedulers.jl schedules are stateless (they don't store their iteration state). If you want a _stateful_ schedule, you can use `ParameterSchedulers.Stateful`:
```julia
using ParameterSchedulers: Stateful, next!
schedule = Stateful(Cos(λ0 = 1e-4, λ1 = 1e-2, period = 10))
for epoch in 1:100
Flux.adjust!(opt_state, next!(schedule))
# your training code here
end
```
Finally, a scheduling function can be incorporated into the optimser's state, advanced at each gradient update step, and possibly passed to the `train!` function. See [this section](https://fluxml.ai/ParameterSchedulers.jl/stable/tutorials/optimizers/#Working-with-Flux-optimizers) of ParameterSchedulers.jl documentation for more details.
ParameterSchedulers.jl allows for many more scheduling policies including arbitrary functions, looping any function with a given period, or sequences of many schedules. See the [ParameterSchedulers.jl documentation](https://fluxml.ai/ParameterSchedulers.jl/stable) for more info.
## Freezing layer parameters
To completely disable training of some part of the model, use [`freeze!`](@ref Flux.freeze!).
This is a temporary modification, reversed by `thaw!`:
```julia
Flux.freeze!(opt_state.layers.enc)
# Now training won't update parameters in bimodel.layers.enc
train!(loss, bimodel, data, opt_state)
# Un-freeze the entire model:
Flux.thaw!(opt_state)
```
While `adjust!` and `freeze!`/`thaw!` make temporary modifications to the optimiser state,
permanently removing some fields of a new layer type from training is usually done
when defining the layer, by calling for example [`@layer`](@ref Flux.@layer)` NewLayer trainable=(weight,)`.
---
### Src/Guide/Gpu
# GPU Support
Most work on neural networks involves the use of GPUs, as they can typically perform the required computation much faster.
This page describes how Flux co-operates with various other packages, which talk to GPU hardware.
For those in a hurry, see the [quickstart](@ref man-quickstart) page. Or do `using CUDA` and then call `gpu` on both the model and the data.
## Basic GPU use: from `Array` to `CuArray`
Julia's GPU packages work with special array types, in place of the built-in `Array`.
The most used is `CuArray` provided by [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl), for GPUs made by NVIDIA.
That package provides a function `cu` which converts an ordinary `Array` (living in CPu memory) to a `CuArray` (living in GPU memory).
Functions like `*` and broadcasting specialise so that, when given `CuArray`s, all the computation happens on the GPU:
```julia
W = randn(3, 4) # some weights, on CPU: 3×4 Array{Float64, 2}
x = randn(4) # fake data
y = tanh.(W * x) # computation on the CPU
using CUDA
cu(W) isa CuArray{Float32}
(cW, cx) = (W, x) |> cu # move both to GPU
cy = tanh.(cW * cx) # computation on the GPU
```
Notice that `cu` doesn't only move arrays, it also recurses into many structures, such as the tuple `(W, x)` above.
(Notice also that it converts Julia's default `Float64` numbers to `Float32`, as this is what most GPUs support efficiently -- it calls itself "opinionated". Flux defaults to `Float32` in all cases.)
To use CUDA with Flux, you can simply use `cu` to move both the model, and the data.
It will create a copy of the Flux model, with all of its parameter arrays moved to the GPU:
```julia
using Pkg; Pkg.add(["CUDA", "cuDNN"]) # do this once
using Flux, CUDA
CUDA.allowscalar(false) # recommended
model = Dense(W, true, tanh) # wrap the same matrix W in a Flux layer
model(x) ≈ y # same result, still on CPU
c_model = cu(model) # move all the arrays within model to the GPU
c_model(cx) # computation on the GPU
```
Notice that you need `using CUDA` (every time) but also `] add cuDNN` (once, when installing packages).
This is a quirk of how these packages are set up.
(The [`cuDNN.jl`](https://github.com/JuliaGPU/CUDA.jl/tree/master/lib/cudnn) sub-package handles operations such as convolutions, called by Flux via [NNlib.jl](https://github.com/FluxML/NNlib.jl).)
Flux's `gradient`, and training functions like `setup`, `update!`, and `train!`, are all equally happy to accept GPU arrays and GPU models, and then perform all computations on the GPU.
It is recommended that you move the model to the GPU before calling `setup`.
```julia
grads = Flux.gradient((f,x) -> sum(abs2, f(x)), model, x) # on CPU
c_grads = Flux.gradient((f,x) -> sum(abs2, f(x)), c_model, cx) # same result, all on GPU
c_opt = Flux.setup(Adam(), c_model) # setup optimiser after moving model to GPU
Flux.update!(c_opt, c_model, c_grads[1]) # mutates c_model but not model
```
To move arrays and other objects back to the CPU, Flux provides a function `cpu`.
This is recommended when saving models, `Flux.state(c_model |> cpu)`, see below.
```julia
cpu(cW) isa Array{Float32, 2}
model2 = cpu(c_model) # copy model back to CPU
model2(x)
```
!!! compat "Flux ≤ 0.13"
Old versions of Flux automatically loaded CUDA.jl to provide GPU support. Starting from Flux v0.14, it has to be loaded separately. Julia's [package extensions](https://pkgdocs.julialang.org/v1/creating-packages/#Conditional-loading-of-code-in-packages-(Extensions)) allow Flux to automatically load some GPU-specific code when needed.
## Other GPU packages for AMD & Apple
Non-NVIDIA graphics cards are supported by other packages. Each provides its own function which behaves like `cu`.
AMD GPU support provided by [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl), on systems with ROCm and MIOpen installed.
This package has a function `roc` which converts `Array` to `ROCArray`:
```julia
using Flux, AMDGPU
AMDGPU.allowscalar(false)
r_model = roc(model)
r_model(roc(x))
Flux.gradient((f,x) -> sum(abs2, f(x)), r_model, roc(x))
```
Experimental support for Apple devices with M-series chips is provided by [Metal.jl](https://github.com/JuliaGPU/Metal.jl). This has a function [`mtl`](https://metal.juliagpu.org/stable/api/array/#Metal.mtl) which works like `cu`, converting `Array` to `MtlArray`:
```julia
using Flux, Metal
Metal.allowscalar(false)
m_model = mtl(model)
m_y = m_model(mtl(x))
Flux.gradient((f,x) -> sum(abs2, f(x)), m_model, mtl(x))
```
!!! danger "Experimental"
Metal support in Flux is experimental and many features are not yet available.
AMD support is improving, but likely to have more rough edges than CUDA.
If you want your model to work with any brand of GPU, or none, then you may not wish to write `cu` everywhere.
One simple way to be generic is, at the top of the file, to un-comment one of several lines which import a package and assign its "adaptor" to the same name:
```julia
using CUDA: cu as device # after this, `device === cu`
# using AMDGPU: roc as device
# device = identity # do-nothing, for CPU
using Flux
model = Chain(...) |> device
```
!!! note "Adapt.jl"
The functions `cu`, `mtl`, `roc` all use [Adapt.jl](https://github.com/JuliaGPU/Adapt.jl), to work within various wrappers.
The reason they work on Flux models is that `Flux.@layer Layer` defines methods of `Adapt.adapt_structure(to, lay::Layer)`.
## Automatic GPU choice with `gpu` and `gpu_device`
Flux also provides a more automatic way of choosing which GPU (or none) to use. This is the function `gpu`:
* By default it does nothing.
* If the package CUDA is loaded, and `CUDA.functional() === true`, then it behaves like `cu`.
* If the package AMDGPU is loaded, and `AMDGPU.functional() === true`, then it behaves like `roc`.
* If the package Metal is loaded, and `Metal.functional() === true`, then it behaves like `mtl`.
* If two different GPU packages are loaded, the first one takes priority.
For the most part, this means that a script which says `model |> gpu` and `data |> gpu` will just work.
It should always run, and if a GPU package is loaded (and finds the correct hardware) then that will be used.
The function `gpu` uses a lower-level function called [`gpu_device`](@ref) from MLDataDevices.jl,
which checks what to do and then returns some device object. In fact, the entire implementation is just this:
```julia
gpu(x) = gpu_device()(x)
cpu(x) = cpu_device()(x)
```
Automatic backend selection through `gpu` is not type-stable. That doesn't matter if you do it once, or once per large batch -- it costs a few microseconds. But it might matter if you do it within some loop.
To avoid this, you can first obtain a "device object" with `device = gpu_device()`, once, and then use this as the function to transfer data. Something like this:
```julia
to_device = gpu_device()
gpu_model = model |> to_device
for epoch in 1:num_epochs
for (x, y) in dataloader
x_gpu, y_gpu = (x, y) |> to_device
# training code...
```
Finally, setting a backend prefence with [`gpu_backend!`](@ref) gives type stability to the whole pipeline.
## Transferring Training Data
In order to train the model using the GPU both model and the training data have to be transferred to GPU memory. Moving the data can be done in two different ways:
1. Iterating over the batches in a [`DataLoader`](@ref) object transferring each one of the training batches at a time to the GPU. This is recommended for large datasets. Done by hand, it might look like this:
```julia
train_loader = Flux.DataLoader((X, Y), batchsize=64, shuffle=true)
# ... model definition, optimiser setup
for epoch in 1:epochs
for (x_cpu, y_cpu) in train_loader
x = gpu(x_cpu)
y = gpu(y_cpu)
grads = gradient(m -> loss(m, x, y), model)
Flux.update!(opt_state, model, grads[1])
end
end
```
Rather than write this out every time, you can just call `gpu(::DataLoader)`:
```julia
gpu_train_loader = Flux.DataLoader((X, Y), batchsize=64, shuffle=true) |> gpu
# ... model definition, optimiser setup
for epoch in 1:epochs
for (x, y) in gpu_train_loader
grads = gradient(m -> loss(m, x, y), model)
Flux.update!(opt_state, model, grads[1])
end
end
```
This is equivalent to `DataLoader(MLUtils.mapobs(gpu, (X, Y)); keywords...)`.
Something similar can also be done with [`CUDA.CuIterator`](https://cuda.juliagpu.org/stable/usage/memory/#Batching-iterator), `gpu_train_loader = CUDA.CuIterator(train_loader)`. However, this only works with a limited number of data types: `first(train_loader)` should be a tuple (or `NamedTuple`) of arrays.
2. Transferring all training data to the GPU at once before creating the `DataLoader`. This is usually performed for smaller datasets which are sure to fit in the available GPU memory.
```julia
gpu_train_loader = Flux.DataLoader((X, Y) |> gpu, batchsize = 32)
# ...
for epoch in 1:epochs
for (x, y) in gpu_train_loader
# ...
```
Here `(X, Y) |> gpu` applies [`gpu`](@ref) to both arrays, as it recurses into structures.
## Saving GPU-Trained Models
After the training process is done, we must always transfer the trained model back to the CPU memory before serializing or saving to disk. This can be done with `cpu`:
```julia
model = cpu(model) # or model = model |> cpu
```
and then
```julia
using BSON
# ...
BSON.@save "./path/to/trained_model.bson" model
# in this approach the cpu-transferred model (referenced by the variable `model`)
# only exists inside the `let` statement
let model = cpu(model)
# ...
BSON.@save "./path/to/trained_model.bson" model
end
# is equivalent to the above, but uses `key=value` storing directive from BSON.jl
BSON.@save "./path/to/trained_model.bson" model = cpu(model)
```
The reason behind this is that models trained in the GPU but not transferred to the CPU memory scope will expect `CuArray`s as input. In other words, Flux models expect input data coming from the same kind device in which they were trained on.
In controlled scenarios in which the data fed to the loaded models is guaranteed to be in the GPU there's no need to transfer them back to CPU memory scope, however in production environments, where artifacts are shared among different processes, equipments or configurations, there is no guarantee that the CUDA.jl package will be available for the process performing inference on the model loaded from the disk.
## Disabling CUDA or choosing which GPUs are visible to Flux
Sometimes it is required to control which GPUs are visible to `julia` on a system with multiple GPUs or disable GPUs entirely. This can be achieved with an environment variable `CUDA_VISIBLE_DEVICES`.
To disable all devices:
```
$ export CUDA_VISIBLE_DEVICES='-1'
```
To select specific devices by device id:
```
$ export CUDA_VISIBLE_DEVICES='0,1'
```
More information for conditional use of GPUs in CUDA.jl can be found in its [documentation](https://cuda.juliagpu.org/stable/installation/conditional/#Conditional-use), and information about the specific use of the variable is described in the [Nvidia CUDA blog post](https://developer.nvidia.com/blog/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/).
## Data movement across GPU devices
Flux also supports getting handles to specific GPU devices, and transferring models from one GPU device to another GPU device from the same backend. Let's try it out for NVIDIA GPUs. First, we list all the available devices:
```julia-repl
julia> using Flux, CUDA;
julia> CUDA.devices()
CUDA.DeviceIterator() for 3 devices:
0. NVIDIA TITAN RTX
1. NVIDIA TITAN RTX
2. NVIDIA TITAN RTX
```
Then, let's select the device with id `0`:
```julia-repl
julia> device0 = gpu_device(1)
(::CUDADevice{CuDevice}) (generic function with 4 methods)
julia> device0.device
CuDevice(0): NVIDIA TITAN RTX
```
Notice that indexing starts from `0` in the `CUDA.devices()` output, but `gpu_device!` expects the device id starting from `1`.
Then, let's move a simple dense layer to the GPU represented by `device0`:
```julia-repl
julia> dense_model = Dense(2 => 3)
Dense(2 => 3) # 9 parameters
julia> dense_model = dense_model |> device0;
julia> dense_model.weight
3×2 CuArray{Float32, 2, CUDA.DeviceMemory}:
-0.142062 -0.131455
-0.828134 -1.06552
0.608595 -1.05375
julia> CUDA.device(dense_model.weight) # check the GPU to which dense_model is attached
CuDevice(0): NVIDIA TITAN RTX
```
Next, we'll get a handle to the device with id `1`, and move `dense_model` to that device:
```julia-repl
julia> device1 = gpu_device(2)
(::CUDADevice{CuDevice}) (generic function with 4 methods)
julia> dense_model = dense_model |> device1; # don't directly print the model; see warning below
julia> CUDA.device(dense_model.weight)
CuDevice(1): NVIDIA TITAN RTX
```
Due to a limitation in `Metal.jl`, currently this kind of data movement across devices is only supported for `CUDA` and `AMDGPU` backends.
## Distributed data parallel training
!!! danger "Experimental"
Distributed support is experimental and could change in the future.
Flux supports now distributed data parallel training with `DistributedUtils` module.
If you want to run your code on multiple GPUs, you have to install `MPI.jl` (see [docs](https://juliaparallel.org/MPI.jl/stable/usage/) for more info).
```julia-repl
julia> using MPI
julia> MPI.install_mpiexecjl()
```
Now you can run your code with `mpiexecjl --project=. -n julia .jl` from CLI.
You can use either the `MPIBackend` or `NCCLBackend`, the latter only if also `NCCL.jl` is loaded. First, initialize a backend with `DistributedUtils.initialize`, e.g.
```julia-repl
julia> using Flux, MPI, NCCL, CUDA
julia> CUDA.allowscalar(false)
julia> DistributedUtils.initialize(NCCLBackend)
julia> backend = DistributedUtils.get_distributed_backend(NCCLBackend)
NCCLBackend{Communicator, MPIBackend{MPI.Comm}}(Communicator(Ptr{NCCL.LibNCCL.ncclComm} @0x000000000607a660), MPIBackend{MPI.Comm}(MPI.Comm(1140850688)))
```
Pass your model, as well as any data to GPU device.
```julia-repl
julia> model = Chain(Dense(1 => 256, tanh), Dense(256 => 1)) |> gpu
Chain(
Dense(1 => 256, tanh), # 512 parameters
Dense(256 => 1), # 257 parameters
) # Total: 4 arrays, 769 parameters, 744 bytes.
julia> x = rand(Float32, 1, 16) |> gpu
1×16 CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}:
0.239324 0.331029 0.924996 0.55593 0.853093 0.874513 0.810269 0.935858 0.477176 0.564591 0.678907 0.729682 0.96809 0.115833 0.66191 0.75822
julia> y = x .^ 3
1×16 CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}:
0.0137076 0.0362744 0.791443 0.171815 0.620854 0.668804 0.53197 0.819654 0.108651 0.179971 0.312918 0.388508 0.907292 0.00155418 0.29 0.435899
```
In this case, we are training on a total of `16 * number of processes` samples. You can also use `DistributedUtils.DistributedDataContainer` to split the data uniformly across processes (or do it manually).
```julia-repl
julia> data = DistributedUtils.DistributedDataContainer(backend, x)
Flux.DistributedUtils.DistributedDataContainer(Float32[0.23932439 0.33102947 … 0.66191036 0.75822026], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
```
You have to wrap your model in `DistributedUtils.FluxDistributedModel` and synchronize it (broadcast across all processes):
```julia-repl
julia> model = DistributedUtils.synchronize!!(backend, DistributedUtils.FluxDistributedModel(model); root=0)
Chain(
Dense(1 => 256, tanh), # 512 parameters
Dense(256 => 1), # 257 parameters
) # Total: 4 arrays, 769 parameters, 744 bytes.
```
Time to set up an optimizer by using `DistributedUtils.DistributedOptimizer` and synchronize it as well.
```
/* Detailed source-code truncated for AI context efficiency. */
```
Now you can define loss and train the model.
```julia-repl
julia> loss(model) = mean((model(x) .- y).^2)
loss (generic function with 1 method)
julia> for epoch in 1:100
global model, st_opt
l, grad = Zygote.withgradient(loss, model)
println("Epoch $epoch: Loss $l")
st_opt, model = Optimisers.update(st_opt, model, grad[1])
end
Epoch 1: Loss 0.011638729
Epoch 2: Loss 0.0116432225
Epoch 3: Loss 0.012763695
...
```
Remember that in order to run it on multiple GPUs you have to run from CLI `mpiexecjl --project=. -n julia .jl`,
where `` is the number of processes that you want to use. The number of processes usually corresponds to the number of gpus.
By default `MPI.jl` MPI installation is CUDA-unaware so if you want to run it in CUDA-aware mode, read more [here](https://juliaparallel.org/MPI.jl/stable/usage/#CUDA-aware-MPI-support) on custom installation and rebuilding `MPI.jl`.
Then test if your MPI is CUDA-aware by
```julia-repl
julia> import Pkg
julia> Pkg.test("MPI"; test_args=["--backend=CUDA"])
```
If it is, set your local preference as below
```julia-repl
julia> using Preferences
julia> set_preferences!("Flux", "FluxDistributedMPICUDAAware" => true)
```
!!! warning "Known shortcomings"
We don't run CUDA-aware tests so you're running it at own risk.
## Checking GPU Availability
By default, Flux will run the checks on your system to see if it can support GPU functionality. You can check if Flux identified a valid GPU setup by typing the following:
```julia-repl
julia> using CUDA
julia> CUDA.functional()
true
```
For AMD GPU:
```julia-repl
julia> using AMDGPU
julia> AMDGPU.functional()
true
julia> AMDGPU.functional(:MIOpen)
true
```
For Metal GPU:
```julia-repl
julia> using Metal
julia> Metal.functional()
true
```
---
### Src/Guide/Performance
# [Performance Tips](@id man-performance-tips)
All the usual [Julia performance tips apply](https://docs.julialang.org/en/v1/manual/performance-tips/).
As always [profiling your code](https://docs.julialang.org/en/v1/manual/profile/#Profiling-1) is generally a useful way of finding bottlenecks.
Below follow some Flux specific tips/reminders.
## Don't use more precision than you need
Flux works great with all kinds of number types.
But often you do not need to be working with say `Float64` (let alone `BigFloat`).
Switching to `Float32` can give you a significant speed up,
not because the operations are faster, but because the memory usage is halved.
Which means allocations occur much faster.
And you use less memory.
## Preserve inputs' types
Not only should your activation and loss functions be [type-stable](https://docs.julialang.org/en/v1/manual/performance-tips/#Write-%22type-stable%22-functions-1),
they should also preserve the type of their inputs.
A very artificial example using an activation function like
```julia
my_tanh(x) = Float64(tanh(x))
```
will result in performance on `Float32` input orders of magnitude slower than the normal `tanh` would,
because it results in having to use slow mixed type multiplication in the dense layers.
Similar situations can occur in the loss function during backpropagation.
Which means if you change your data say from `Float64` to `Float32` (which should give a speedup: see above),
you will see a large slow-down.
This can occur sneakily, because you can cause type-promotion by interacting with a numeric literals.
E.g. the following will have run into the same problem as above:
```julia
leaky_tanh(x) = 0.01*x + tanh(x)
```
While one could change the activation function (e.g. to use `0.01f0*x`), the idiomatic (and safe way) to avoid type casts whenever inputs changes is to use `oftype`:
```julia
leaky_tanh(x) = oftype(x/1, 0.01)*x + tanh(x)
```
## Evaluate batches as matrices of features
While it can sometimes be tempting to process your observations (feature vectors) one at a time
e.g.
```julia
function loss_total(xs::AbstractVector{<:Vector}, ys::AbstractVector{<:Vector})
sum(zip(xs, ys)) do (x, y_target)
y_pred = model(x) # evaluate the model
return loss(y_pred, y_target)
end
end
```
It is much faster to concatenate them into a matrix,
as this will hit BLAS matrix-matrix multiplication, which is much faster than the equivalent sequence of matrix-vector multiplications.
The improvement is enough that it is worthwhile allocating new memory to store them contiguously.
```julia
x_batch = reduce(hcat, xs)
y_batch = reduce(hcat, ys)
...
function loss_total(x_batch::Matrix, y_batch::Matrix)
y_preds = model(x_batch)
sum(loss.(y_preds, y_batch))
end
```
When doing this kind of concatenation use `reduce(hcat, xs)` rather than `hcat(xs...)`.
This will avoid the splatting penalty, and will hit the optimised `reduce` method.
## Be aware of GPU memory inefficiencies
Currently, GPU memory is not handled as well as system memory. If your training loop is allocating significantly on the GPU, you can quickly fill your GPU memory and the piecemeal reclamation and shuffling of data between GPU and system memory can become extremely slow. If profiling shows that a significant portion of time is spent in the `gpu` function and your data sizes are not large, this may be the cause. Running an incremental garbage collection manually (`GC.gc(false)`) at regular intervals can keep your GPU memory free and responsive. See other tips for CUDA memory management [here](https://cuda.juliagpu.org/stable/usage/memory/).
## Compile your training with Reactant
The tips above tune *eager* execution. For the largest speedups, compile your model with [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl): it traces your code and lowers it — through MLIR and XLA — into a single fused, kernel-optimised executable that is cached and reused across calls. For training, the forward pass, the Enzyme reverse pass and the optimiser update are compiled *together*, which removes per-operation dispatch and kernel-launch overhead and lets the compiler fuse operations and plan buffers across the whole step.
You do not have to invoke the compiler yourself: when the model lives on a Reactant device, [`trainstep!`](@ref Flux.Train.trainstep!) — and [`train!`](@ref Flux.train!), which is built on it — compile and cache the fused step automatically.
```julia
using Flux, Reactant
dev = reactant_device()
model = model |> dev
opt_state = Flux.setup(Adam(1f-3), model) # set up the optimiser after moving to the device
loader = DataLoader((X, Y); batchsize=32) |> dev # move batches to the device during iteration
loss(m, x, y) = Flux.logitcrossentropy(m(x), y)
trainmode!(model)
for (x, y) in loader
Flux.trainstep!(loss, model, (x, y), opt_state) # compiled on the first call, reused after
end
```
The cached executable is keyed on the batch shape, so keeping batch sizes fixed avoids recompilation (a smaller final batch simply compiles one extra executable). The first call pays a one-time compilation cost that can be large, so this pays off over many steps rather than a handful. See [Compiling Flux with Reactant](reactant.md) for the full guide, including inference and manual `@compile`.
---
### Src/Guide/Reactant
# Compiling Flux with Reactant
[Reactant.jl](https://github.com/EnzymeAD/Reactant.jl) traces your Julia code and compiles it — through [MLIR](https://mlir.llvm.org/) and [XLA](https://openxla.org/xla) — into a single optimised executable that runs on CPU, NVIDIA/AMD GPUs, or TPUs. For Flux this means the whole forward pass (and, for training, the [Enzyme](https://enzyme.mit.edu/) reverse pass and the optimiser update) is fused, kernel-optimised, and reused across calls, which is often substantially faster and more memory-efficient than eager execution.
This guide builds up in three steps:
1. A **manually compiled** example, so you can see exactly what Reactant does.
2. The [`trainstep!`](@ref Flux.Train.trainstep!) API, which compiles and caches a single training step for you.
3. [`train!`](@ref Flux.train!), the full training loop, which is built on top of `trainstep!`.
## Installation
Reactant is a normal dependency — add it and load it alongside Flux:
```julia
using Pkg; Pkg.add("Reactant") # do this once
using Flux, Reactant
```
Reactant automatically selects a GPU backend if one is available, falling back to the CPU otherwise. See the [Reactant GPU configuration docs](https://enzymead.github.io/Reactant.jl/dev/api/config#GPU-Configuration) for how to control this.
## Moving data and models to the device
Just like `gpu`/`cpu`, Flux provides a device object that moves a model (or any nested structure of arrays) onto the Reactant device, converting its arrays into Reactant's `ConcreteRArray`s:
```julia
using Flux, Reactant
dev = reactant_device() # a Reactant device object; call once and reuse
model = Chain(Dense(4 => 8, tanh), Dense(8 => 2))
x = randn(Float32, 4, 16)
model_re = model |> dev # a copy of the model with Reactant arrays
x_re = x |> dev
```
As with the GPU adaptors, `dev` recurses into structures, so `(x, y) |> dev` moves both, and it uses `Float32` by default. Everything you feed to a compiled function must already live on the device.
## 1. A manually compiled example
Reactant does not run your code eagerly. Instead you *compile* a function once for a given set of input shapes and types, and then call the returned executable. The two entry points are:
- `Reactant.@compile f(args...)` — trace and compile `f`, returning a callable executable. Call it later with device-resident arguments of the same shapes.
- `Reactant.@jit f(args...)` — compile *and* immediately run, a convenient shortcut for one-off calls.
### Compiling the forward pass (inference)
```julia
using Flux, Reactant
dev = reactant_device()
model = Chain(Dense(4 => 8, tanh), Dense(8 => 2)) |> dev
x = randn(Float32, 4, 16) |> dev
# Compile once...
forward = Reactant.@compile model(x)
# ...then call the compiled executable (fast, no recompilation):
y = forward(x) # a Reactant array on the device
y_host = y |> cpu # move the result back to the host
# `@jit` compiles and runs in one go — handy for a one-off evaluation:
y2 = Reactant.@jit model(x)
```
Note the call is `forward(x)`, **not** `forward(model, x)`: in `@compile model(x)` the model is the *callee*, so Reactant captures it (and its parameter arrays) inside the compiled executable, and you pass only the remaining arguments. Because the model is captured by reference, mutating its parameters in place — as training does — is reflected the next time you call `forward`. (Contrast this with the training-step examples below, where the model sits in an *argument* position, e.g. `@compile my_step!(loss, model, x, y, opt_state)`, and so must be passed at call time.)
A compiled executable is specialised to the **shapes** of its inputs. Calling `forward` with a differently-shaped `x` (e.g. a smaller final batch) requires compiling a separate executable for that shape.
!!! note "trainmode / testmode"
Layers such as `Dropout` and `BatchNorm` behave differently during training and inference. Call `testmode!(model)` before compiling an inference function, and `trainmode!(model)` before a training step, exactly as you would without Reactant.
### Compiling a training step manually
To differentiate under Reactant, use Flux's [`withgradient`](@ref Flux.withgradient) with the [`AutoEnzyme`](@ref) backend — Reactant relies on Enzyme for AD. You can compile the value-and-gradient computation, or a whole step that also updates the model in place:
```julia
using Flux, Reactant, Optimisers
dev = reactant_device()
model = Chain(Dense(4 => 8, tanh), Dense(8 => 2)) |> dev
x, y = randn(Float32, 4, 16) |> dev, randn(Float32, 2, 16) |> dev
loss(m, x, y) = Flux.mse(m(x), y)
opt_state = Flux.setup(Adam(1f-2), model)
# A plain Julia function describing one optimisation step. It differentiates the loss
# with Enzyme, updates the model and optimiser state in place, and returns the loss.
function my_step!(loss, model, x, y, opt_state)
l, grads = Flux.withgradient(m -> loss(m, x, y), AutoEnzyme(), model)
Optimisers.update!(opt_state, model, grads[1])
return l
end
trainmode!(model)
step! = Reactant.@compile my_step!(loss, model, x, y, opt_state) # compile once
for epoch in 1:100
l = step!(loss, model, x, y, opt_state) # reuse the executable each epoch
@info "epoch $epoch" loss=Reactant.to_number(l)
end
```
Reactant traces the mutation of `model` and `opt_state` and fuses the whole step. Writing this by hand gives you full control, but it is boilerplate that Flux can handle for you — which is what the next section is about.
## 2. The `trainstep!` API
[`trainstep!`](@ref Flux.Train.trainstep!) performs exactly the step above — differentiate the loss, update `model` and `opt_state` in place, return the loss — and when the model lives on a Reactant device it **compiles and caches the fused step automatically**. You do not write `@compile` yourself, and repeated calls with the same model, optimiser, loss and batch shape reuse the cached executable.
```julia
using Flux, Reactant
dev = reactant_device()
model = Chain(Dense(4 => 8, tanh), Dense(8 => 2)) |> dev
x, y = randn(Float32, 4, 16) |> dev, randn(Float32, 2, 16) |> dev
loss(m, x, y) = Flux.mse(m(x), y)
# Move the model to the device *before* `setup`, so the optimiser is set up for Reactant arrays.
opt_state = Flux.setup(Adam(1f-2), model)
trainmode!(model)
for epoch in 1:100
l = Flux.trainstep!(loss, model, (x, y), opt_state) # returns the host-side loss scalar
@info "epoch $epoch" loss=l
end
```
The batch is passed as a tuple `(x, y)` and spliced into the loss as `loss(model, x, y)`. The returned loss is read back to the host as an ordinary number.
Use [`trainstep_withgradient!`](@ref Flux.Train.trainstep_withgradient!) if you also need the gradient — it returns `(loss, grad)`, with the gradient left on the device. (Returning the gradient makes it an output of the compiled step and raises peak memory, so prefer `trainstep!` when you don't need it.)
### Auxiliary loss outputs
Like `withgradient`, the loss may return auxiliary data alongside the scalar loss: return a `Tuple` or `NamedTuple` whose first element is the loss, and the gradient is taken of the loss alone while the whole value is returned. This is convenient for logging a metric computed in the forward pass — and, on Reactant, the metric is computed on-device as part of the same compiled forward:
```julia
function loss(m, x, y)
ŷ = m(x)
Flux.mse(ŷ, y), (; acc = mean(onecold(ŷ) .== onecold(y)))
end
l, stats = Flux.trainstep!(loss, model, (x, y), opt_state) # l == (loss, stats); stats.acc read to host
```
## 3. The `train!` API
[`train!`](@ref Flux.train!) is a loop over the data built on top of `trainstep!`, so on a Reactant device it inherits the same automatic compile-and-cache. Move the model to the device and call `setup` first, and provide **device-resident** data — `train!` rejects host arrays on the Reactant path.
```julia
using Flux, Reactant
dev = reactant_device()
model = Chain(Dense(4 => 8, tanh), Dense(8 => 2)) |> dev
opt_state = Flux.setup(Adam(1f-2), model)
loss(m, x, y) = Flux.mse(m(x), y)
# Move every batch to the device. In a real loop use a DataLoader wrapped with the device,
# so each batch is moved lazily and the previous one is freed:
# train_loader = DataLoader((X, Y); batchsize=32) |> dev
X, Y = randn(Float32, 4, 128), randn(Float32, 2, 128)
data = [(X[:, i:i+15], Y[:, i:i+15]) |> dev for i in 1:16:128]
Flux.train!(loss, model, data, opt_state)
```
`train!` runs a single step per batch, shows a progress bar, and stops with a `DomainError` if the loss becomes non-finite. Because the compiled step is cached and keyed on the batch shape, multi-epoch training does not recompile, and a smaller final batch simply compiles one additional executable that is then reused.
## Tips and gotchas
- **`setup` after moving to the device.** Always `model |> reactant_device()` *before* `Flux.setup`, so the optimiser state is created from Reactant arrays (this keeps stateful rules like `Adam`'s bias-correction term on the device).
- **Device-resident data.** Every input to a compiled step must already live on the device. Move batches with `|> reactant_device()`; a common pattern is to wrap a `DataLoader` with the device so batches are moved lazily.
- **Compilation is keyed on shape.** The first call for each distinct batch shape compiles; subsequent calls reuse the executable. Keeping batch sizes fixed (or accepting one extra compile for a smaller final batch) avoids repeated compilation.
- **Watch the compile cache.** Flux caches one executable per distinct `(model, optimiser, loss, batch-shape)`. If you rebuild the model or optimiser every iteration, each iteration compiles a fresh step — Flux warns once the cache grows past a handful of entries. Entries are freed automatically when their model is garbage-collected.
- **Reading results back.** `trainstep!`/`train!` return host-side numbers already. For a manually compiled function, move device arrays back with `cpu` (or `Reactant.to_number` for a scalar).
- **AD backend.** On a Reactant device training differentiates with Enzyme by default. Passing an explicit `adtype` such as `AutoZygote()` or `AutoMooncake()` compilation is likely to fail. `Duplicated` models are not supported here — pass the plain model that already lives on the device.
See also the [`resnet_tinyimagenet` example](https://github.com/FluxML/Flux.jl/tree/master/examples/resnet_tinyimagenet) for a complete Reactant training script, including compiling a separate evaluation executable per batch shape.
---
### Src/Guide/Saving
# Saving and Loading Models
You may wish to save models so that they can be loaded and run in a later
session. Flux provides a number of ways to do this.
The recommended way, which is the most robust one for long term storage,
is to use [`Flux.state`](@ref) in combination with a serialization format like
[JLD2.jl](https://juliaio.github.io/JLD2.jl/dev/) or
[BSON.jl](https://github.com/JuliaIO/BSON.jl).
Save a model:
```jldoctest saving
julia> using Flux
julia> struct MyModel
net
end
julia> Flux.@layer MyModel
julia> MyModel() = MyModel(Chain(Dense(10 => 5, relu), Dense(5 => 2)));
julia> model = MyModel()
MyModel(
Chain(
Dense(10 => 5, relu), # 55 parameters
Dense(5 => 2), # 12 parameters
),
) # Total: 4 arrays, 67 parameters, 484 bytes.
julia> model_state = Flux.state(model);
julia> using JLD2
julia> jldsave("mymodel.jld2"; model_state)
```
Load it again in a new session using [`Flux.loadmodel!`](@ref):
```jldoctest saving
julia> using Flux, JLD2
julia> model_state = JLD2.load("mymodel.jld2", "model_state");
julia> model = MyModel(); # MyModel definition must be available
julia> Flux.loadmodel!(model, model_state);
```
!!! note
If a saved model's parameters are stored on the GPU, the model will not load
later on if there is no GPU support available. It's best to [move your model
to the CPU](gpu.md) with `cpu(model)` before saving it.
## Checkpointing
In longer training runs it's a good idea to periodically save your model, so that you can resume if training is interrupted (for example, if there's a power cut).
```jldoctest saving
julia> using Flux: throttle
julia> using JLD2
julia> m = Chain(Dense(10 => 5, relu), Dense(5 => 2))
Chain(
Dense(10 => 5, relu), # 55 parameters
Dense(5 => 2), # 12 parameters
) # Total: 4 arrays, 67 parameters, 476 bytes.
julia> for epoch in 1:10
# ... train model ...
jldsave("model-checkpoint.jld2", model_state = Flux.state(m))
end;
```
This will update the `"model-checkpoint.jld2"` every epoch.
You can get more advanced by saving a series of models throughout training, for example
```julia
jldsave("model-$(now()).jld2", model_state = Flux.state(m))
```
will produce a series of models like `"model-2018-03-06T02:57:10.41.jld2"`. You
could also store the current test set loss, so that it's easy to (for example)
revert to an older copy of the model if it starts to overfit.
```julia
jldsave("model-$(now()).jld2", model_state = Flux.state(m), loss = testloss())
```
Note that to resume a model's training, you might need to restore other stateful parts of your training loop. Possible examples are the optimiser state and the randomness used to partition the original data into the training and validation sets.
You can store the optimiser state alongside the model, to resume training
exactly where you left off:
```julia
model = MyModel()
opt_state = Flux.setup(AdamW(), model)
# ... train model ...
model_state = Flux.state(model)
jldsave("checkpoint_epoch=42.jld2"; model_state, opt_state)
```
# Saving Models as Julia Structs
Models are just normal Julia structs, so it's fine to use any Julia storage
format to save the struct as it is instead of saving the state returned by [`Flux.state`](@ref).
[BSON.jl](https://github.com/JuliaIO/BSON.jl) is particularly convenient for this,
since it can also save anonymous functions, which are sometimes part of a model definition.
Save a model:
```jldoctest saving
julia> using Flux
julia> model = Chain(Dense(10 => 5, NNlib.relu), Dense(5 => 2));
julia> using BSON: @save
julia> @save "mymodel.bson" model
```
Load it again in a new session:
```jldoctest saving
julia> using Flux, BSON
julia> BSON.@load "mymodel.bson" model
julia> model
Chain(
Dense(10 => 5, relu), # 55 parameters
Dense(5 => 2), # 12 parameters
) # Total: 4 arrays, 67 parameters, 476 bytes.
```
!!! warning
Saving models this way could lead to compatibility issues across julia versions
and across Flux versions if some of the Flux layers' internals are changed.
It is therefore not recommended for long term storage, use [`Flux.state`](@ref) instead.
---
### Src/Reference/Data/Mldatadevices
```@meta
CurrentModule = MLDataDevices
CollapsedDocStrings = true
```
# Transferring data across devices
Flux relies on the MLDataDevices.jl package to manage devices and transfer data across them. You don't have to explicitly use the package, as Flux re-exports the necessary functions and types.
```@docs
MLDataDevices.cpu_device
MLDataDevices.default_device_rng
MLDataDevices.functional
MLDataDevices.get_device
MLDataDevices.gpu_device
MLDataDevices.gpu_backend!
MLDataDevices.get_device_type
MLDataDevices.isleaf
MLDataDevices.loaded
MLDataDevices.reset_gpu_device!
MLDataDevices.set_device!
MLDataDevices.supported_gpu_backends
MLDataDevices.DeviceIterator
```
---
### Src/Reference/Data/Mlutils
```@meta
CurrentModule = Flux
CollapsedDocStrings = true
```
# Working with Data, using MLUtils.jl
Flux re-exports the `DataLoader` type and utility functions for working with
data from [MLUtils](https://github.com/JuliaML/MLUtils.jl).
## `DataLoader`
The `DataLoader` can be used to create mini-batches of data, in the format [`train!`](@ref Flux.train!) expects.
```@docs
MLUtils.DataLoader
```
## Utility Functions
The utility functions are meant to be used while working with data;
these functions help create inputs for your models or batch your dataset.
```@docs
MLUtils.batch
MLUtils.batchsize
MLUtils.batchseq
MLUtils.batch_sequence
MLUtils.BatchView
MLUtils.chunk
MLUtils.eachobs
MLUtils.fill_like
MLUtils.filterobs
Flux.flatten
MLUtils.flatten
MLCore.getobs
MLCore.getobs!
MLUtils.joinobs
MLUtils.group_counts
MLUtils.group_indices
MLUtils.groupobs
MLUtils.kfolds
MLUtils.leavepout
MLUtils.mapobs
MLCore.numobs
MLUtils.normalise
MLUtils.obsview
MLUtils.ObsView
MLUtils.ones_like
MLUtils.oversample
MLUtils.randobs
MLUtils.rand_like
MLUtils.randn_like
MLUtils.rpad_constant
MLUtils.shuffleobs
MLUtils.splitobs
MLUtils.unbatch
MLUtils.undersample
MLUtils.unsqueeze
MLUtils.unstack
MLUtils.zeros_like
```
---
### Src/Reference/Data/Onehot
```@meta
CollapsedDocStrings = true
```
# One-Hot Encoding with OneHotArrays.jl
It's common to encode categorical variables (like `true`, `false` or `cat`, `dog`) in "one-of-k" or ["one-hot"](https://en.wikipedia.org/wiki/One-hot) form. [OneHotArrays.jl](https://github.com/FluxML/OneHotArrays.jl) provides the `onehot` function to make this easy.
```jldoctest onehot
julia> using OneHotArrays
julia> onehot(:b, [:a, :b, :c])
3-element OneHotVector(::UInt32) with eltype Bool:
⋅
1
⋅
julia> onehot(:c, [:a, :b, :c])
3-element OneHotVector(::UInt32) with eltype Bool:
⋅
⋅
1
```
There is also a `onecold` function, which is an inverse of `onehot`. It can also be given an array of numbers instead of booleans, in which case it performs an `argmax`-like operation, returning the label with the highest corresponding weight.
```jldoctest onehot
julia> onecold(ans, [:a, :b, :c])
:c
julia> onecold([true, false, false], [:a, :b, :c])
:a
julia> onecold([0.3, 0.2, 0.5], [:a, :b, :c])
:c
```
For multiple samples at once, `onehotbatch` creates a batch (matrix) of one-hot vectors, and `onecold` treats matrices as batches.
```jldoctest onehot
julia> using OneHotArrays
julia> onehotbatch([:b, :a, :b], [:a, :b, :c])
3×3 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
⋅ 1 ⋅
1 ⋅ 1
⋅ ⋅ ⋅
julia> onecold(ans, [:a, :b, :c])
3-element Vector{Symbol}:
:b
:a
:b
```
Note that these operations returned `OneHotVector` and `OneHotMatrix` rather than `Array`s. `OneHotVector`s behave like normal vectors but avoid any unnecessary cost compared to using an integer index directly. For example, multiplying a matrix with a one-hot vector simply slices out the relevant row of the matrix under the hood.
## Function listing
```@docs
OneHotArrays.onehot
OneHotArrays.onecold
OneHotArrays.onehotbatch
OneHotArrays.OneHotArray
OneHotArrays.OneHotVector
OneHotArrays.OneHotMatrix
```
---
### Src/Reference/Models/Activation
```@meta
CollapsedDocStrings = true
```
# [Activation Functions from NNlib.jl](@id man-activations)
These non-linearities used between layers of your model are exported by the [NNlib](https://github.com/FluxML/NNlib.jl) package.
Note that, unless otherwise stated, activation functions operate on scalars. To apply them to an array you can call `σ.(xs)`, `relu.(xs)` and so on. Alternatively, they can be passed to a layer like `Dense(784 => 1024, relu)` which will handle this broadcasting.
Functions like [`softmax`](@ref) are sometimes described as activation functions, but not by Flux. They must see all the outputs, and hence cannot be broadcasted. See the next page for details.
## Alphabetical Listing
```@docs
celu
elu
gelu
hardsigmoid
hardswish
hardtanh
leakyrelu
lisht
logcosh
logsigmoid
mish
relu
relu6
rrelu
selu
sigmoid
sigmoid_fast
softplus
softshrink
softsign
swish
tanhshrink
tanh_fast
trelu
```
## One More
Julia's `Base.Math` also provides `tanh`, which can be used as an activation function.
Note that many Flux layers will automatically replace this with [`NNlib.tanh_fast`](@ref) when called, as Base's `tanh` is slow enough to sometimes be a bottleneck.
```julia-repl
julia> using UnicodePlots
julia> lineplot(tanh, -3, 3, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⣀⠤⠔⠒⠒⠉⠉⠉⠉⠉⠉⠉⠉⠉│ tanh(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⡠⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⡰⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⡤⡯⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠎⠁⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠴⠊⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⡤⠤⠔⠒⠉⠁⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
```
---
### Src/Reference/Models/Functors
```@meta
CollapsedDocStrings = true
```
# Recursive transformations from Functors.jl
Flux models are deeply nested structures, and [Functors.jl](https://github.com/FluxML/Functors.jl) provides tools needed to explore such objects, apply functions to the parameters they contain (e.g. for moving them to gpu), and re-build them.
!!! compat "Flux ≤ v0.14"
All layers were previously defined with the `Functors.@functor` macro.
This still works, but it is recommended that you use the new [`Flux.@layer`](@ref Flux.@layer) macro instead.
Both allow [`Flux.setup`](@ref Flux.setup) to see the parameters inside, and [`gpu`](@ref) to move them to the GPU, but [`Flux.@layer`](@ref Flux.@layer) also overloads printing,
and offers a way to define `trainable` at the same time.
!!! compat "Functors v0.5"
With Functors.jl v0.5, which is required by Flux v0.15 and later, every custom type is a functor by default. This means that applying `Flux.@layer` to a type is no longer strictly necessary, but it is still recommended for addictional features like pretty-printing.
`Functors.jl` has its own [notes on basic usage](https://fluxml.ai/Functors.jl/stable/#Basic-Usage-and-Implementation) for more details. Additionally, the [Advanced Model Building and Customisation](@ref man-advanced) page covers the use cases of `Functors` in greater details.
```@docs
Flux.@layer
Functors.@leaf
Functors.@functor
Functors.fmap
Functors.fmap_with_path
Functors.isleaf
Functors.children
Functors.fcollect
Functors.functor
Functors.fmapstructure
Functors.fmapstructure_with_path
Functors.execute
Functors.AbstractWalk
Functors.ExcludeWalk
Functors.CachedWalk
```
## Moving models, or data, to the GPU
Flux provides some convenience functions based on `fmap`. Some ([`f16`](@ref Flux.f16), [`bf16`](@ref Flux.bf16), [`f32`](@ref Flux.f32), [`f64`](@ref Flux.f64)) change the precision of all arrays in a model. Others are used for moving a model to of from GPU memory:
```@docs
cpu
gpu(::Any)
gpu(::Flux.DataLoader)
```
---
### Src/Reference/Models/Layers
```@meta
CollapsedDocStrings = true
```
# [Built-in Layer Types](@id man-layers)
If you started at the beginning of the guide, then you have already met the
basic [`Dense`](@ref) layer, and seen [`Chain`](@ref) for combining layers.
These core layers form the foundation of almost all neural networks.
The `Dense` exemplifies several features:
* It contains an an [activation function](@ref man-activations), which is broadcasted over the output. Because this broadcast can be fused with other operations, doing so is more efficient than applying the activation function separately.
* It take an `init` keyword, which accepts a function acting like `rand`. That is, `init(2,3,4)` should create an array of this size. Flux has [many such functions](@ref man-init-funcs) built-in. All make a CPU array, moved later with [`gpu`](@ref Flux.gpu) if desired.
* The bias vector is always initialised [`Flux.zeros32`](@ref). The keyword `bias=false` will turn this off, i.e. keeping the bias permanently zero.
* It is annotated with [`@layer`](@ref Flux.@layer), which means that [`Flux.setup`](@ref Flux.setup) will see the contents, and [`gpu`](@ref Flux.gpu) will move their arrays to the GPU.
By contrast, `Chain` itself contains no parameters, but connects other layers together.
The section on [dataflow layers](@ref man-dataflow-layers) introduces others like this.
## Fully Connected
```@docs
Dense
Flux.Bilinear
Flux.Scale
```
Perhaps `Scale` isn't quite fully connected, but it may be thought of as `Dense(Diagonal(s.weights), s.bias)`, and LinearAlgebra's `Diagonal` is a matrix which just happens to contain many zeros.
## Convolution Models
These layers are used to build convolutional neural networks (CNNs).
They all expect images in what is called WHCN order: a batch of 32 colour images, each 50 x 50 pixels, will have `size(x) == (50, 50, 3, 32)`. A single grayscale image might instead have `size(x) == (28, 28, 1, 1)`.
Besides images, 2D data, they also work with 1D data, where for instance stereo sound recording with 1000 samples might have `size(x) == (1000, 2, 1)`. They will also work with 3D data, `ndims(x) == 5`, where again the last two dimensions are channel and batch.
To understand how strides and padding work, the article by [Dumoulin & Visin](https://arxiv.org/abs/1603.07285) has great illustrations.
```@docs
Conv
ConvTranspose
CrossCor
DepthwiseConv
```
## MultiHeadAttention
The basic blocks needed to implement [Transformer](https://arxiv.org/abs/1706.03762) architectures. See also the functional counterparts
documented in NNlib's [Attention](@ref) section.
```@docs
MultiHeadAttention
```
### Pooling
These layers are commonly used after a convolution layer, and reduce the size of its output. They have no trainable parameters.
```@docs
AdaptiveMaxPool
MaxPool
GlobalMaxPool
AdaptiveMeanPool
MeanPool
GlobalMeanPool
```
## Upsampling
The opposite of pooling, these layers increase the size of an array. They have no trainable parameters.
```@docs
Upsample
PixelShuffle
```
## Embedding Vectors
These layers accept an index, and return a vector (or several indices, and several vectors). The possible embedding vectors are learned parameters.
```@docs
Flux.Embedding
Flux.EmbeddingBag
```
## [Dataflow Layers, or Containers](@id man-dataflow-layers)
The basic `Chain(F, G, H)` applies the layers it contains in sequence, equivalent to `H ∘ G ∘ F`. Flux has some other layers which contain layers, but connect them up in a more complicated way: `SkipConnection` allows ResNet's residual connection.
```@docs
Chain
Flux.activations
Maxout
SkipConnection
Parallel
PairwiseFusion
```
## Recurrent Models
Much like the core layers above, but can be used to process sequence data (as well as other kinds of structured data).
```@docs
Recurrence
RNNCell
RNN
LSTMCell
LSTM
GRUCell
GRU
GRUv3Cell
GRUv3
Flux.initialstates
```
## Normalisation & Regularisation
These layers don't affect the structure of the network but may improve training times or reduce overfitting. Some of them contain trainable parameters, while others do not.
```@docs
BatchNorm
Dropout
AlphaDropout
LayerNorm
InstanceNorm
GroupNorm
WeightNorm
Flux.remove_weight_norms
Flux.normalise
```
### Test vs. Train
Several normalisation layers behave differently under training and inference (testing). By default, Flux will automatically determine when a layer evaluation is part of training or inference.
!!! warning
This automatic train/test detection works best with Zygote, the default
automatic differentiation package. It may not work with other packages
such as Tracker, Yota, or ForwardDiff.
The functions `Flux.trainmode!` and `Flux.testmode!` let you manually specify which behaviour you want. When called on a model, they will place all layers within the model into the specified mode.
```@docs
testmode!
trainmode!
```
---
### Src/Reference/Models/Losses
```@meta
CollapsedDocStrings = true
```
# [Loss Functions](@id man-losses)
Flux provides a large number of common loss functions used for training machine learning models.
They are grouped together in the `Flux.Losses` module.
Loss functions for supervised learning typically expect as inputs a target `y`, and a prediction `ŷ` from your model.
In Flux's convention, the order of the arguments is the following
```julia
loss(ŷ, y)
```
Most loss functions in Flux have an optional argument `agg`, denoting the type of aggregation performed over the
batch:
```julia
loss(ŷ, y) # defaults to `mean`
loss(ŷ, y, agg=sum) # use `sum` for reduction
loss(ŷ, y, agg=x->sum(x, dims=2)) # partial reduction
loss(ŷ, y, agg=x->mean(w .* x)) # weighted mean
loss(ŷ, y, agg=identity) # no aggregation.
```
## Function listing
```@docs
Flux.Losses.mae
Flux.Losses.mse
Flux.Losses.msle
Flux.Losses.huber_loss
Flux.Losses.label_smoothing
Flux.Losses.crossentropy
Flux.Losses.logitcrossentropy
Flux.Losses.binarycrossentropy
Flux.Losses.logitbinarycrossentropy
Flux.Losses.kldivergence
Flux.Losses.poisson_loss
Flux.Losses.hinge_loss
Flux.Losses.squared_hinge_loss
Flux.Losses.dice_coeff_loss
Flux.Losses.tversky_loss
Flux.Losses.binary_focal_loss
Flux.Losses.focal_loss
Flux.Losses.siamese_contrastive_loss
```
---
### Src/Reference/Models/Nnlib
```@meta
CollapsedDocStrings = true
```
# Neural Network primitives from NNlib.jl
Flux re-exports all of the functions exported by the [NNlib](https://github.com/FluxML/NNlib.jl) package. This includes activation functions, described on [their own page](@ref man-activations). Many of the functions on this page exist primarily as the internal implementation of Flux layer, but can also be used independently.
## Attention
Primitives for the [`MultiHeadAttention`](@ref) layer.
```@docs
NNlib.dot_product_attention
NNlib.dot_product_attention_scores
NNlib.make_causal_mask
```
## Softmax
`Flux`'s [`Flux.logitcrossentropy`](@ref) uses [`NNlib.logsoftmax`](@ref) internally.
```@docs
softmax
logsoftmax
```
## Pooling
`Flux`'s [`AdaptiveMaxPool`](@ref), [`AdaptiveMeanPool`](@ref), [`GlobalMaxPool`](@ref), [`GlobalMeanPool`](@ref),
[`MaxPool`](@ref), and [`MeanPool`](@ref) use [`NNlib.PoolDims`](@ref), [`NNlib.maxpool`](@ref), and [`NNlib.meanpool`](@ref) as their backend.
```@docs
NNlib.PoolDims
NNlib.lpnormpool
NNlib.maxpool
NNlib.meanpool
```
## Padding
```@docs
NNlib.pad_circular
NNlib.pad_constant
NNlib.pad_reflect
NNlib.pad_repeat
NNlib.pad_symmetric
NNlib.pad_zeros
```
## Convolution
`Flux`'s [`Conv`](@ref) and [`CrossCor`](@ref) layers use [`NNlib.DenseConvDims`](@ref) and [`NNlib.conv`](@ref) internally.
```@docs
conv
ConvDims
depthwiseconv
DepthwiseConvDims
DenseConvDims
```
## Dropout
```@docs
NNlib.dropout
NNlib.dropout!
```
## Normalization
`Flux`'s [`BatchNorm`](@ref), [`InstanceNorm`](@ref), [`GroupNorm`](@ref), and [`LayerNorm`](@ref) layers wrap the functional normalization operators below as their backend.
```@docs
NNlib.normalise
NNlib.batchnorm
NNlib.instancenorm
NNlib.groupnorm
NNlib.layernorm
```
## Upsampling
`Flux`'s [`Upsample`](@ref) layer uses [`NNlib.upsample_nearest`](@ref), [`NNlib.upsample_bilinear`](@ref), and [`NNlib.upsample_trilinear`](@ref) as its backend. Additionally, `Flux`'s [`PixelShuffle`](@ref) layer uses [`NNlib.pixel_shuffle`](@ref) as its backend.
```@docs
upsample_nearest
upsample_linear
∇upsample_linear
upsample_bilinear
∇upsample_bilinear
upsample_trilinear
∇upsample_trilinear
pixel_shuffle
```
## Batched Operations
`Flux`'s [`Flux.Bilinear`](@ref) layer uses [`NNlib.batched_mul`](@ref) internally.
```@docs
batched_mul
batched_mul!
batched_adjoint
batched_transpose
batched_vec
```
## Gather and Scatter
`Flux`'s [`Embedding`](@ref) layer uses [`NNlib.gather`](@ref) as its backend.
```@docs
NNlib.gather
NNlib.gather!
NNlib.scatter
NNlib.scatter!
```
## Sampling
```@docs
grid_sample
∇grid_sample
```
## Losses
```@docs
ctc_loss
```
## Miscellaneous
```@docs
logsumexp
NNlib.glu
```
---
### Src/Reference/Training/Callbacks
```@meta
CollapsedDocStrings = true
```
# [Callback Helpers](@id man-callback-helpers)
```@docs
Flux.throttle
```
## Patience Helpers
Flux provides utilities for controlling your training procedure according to some monitored condition and a maximum `patience`. For example, you can use `early_stopping` to stop training when the model is converging or deteriorating, or you can use `plateau` to check if the model is stagnating.
For example, below we create a pseudo-loss function that decreases, bottoms out, and then increases. The early stopping trigger will break the loop before the loss increases too much.
```julia
# create a pseudo-loss that decreases for 4 calls, then starts increasing
# we call this like loss()
loss = let t = 0
() -> begin
t += 1
(t - 4) ^ 2
end
end
# create an early stopping trigger
# returns true when the loss increases for two consecutive steps
es = early_stopping(loss, 2; init_score = 9)
# this will stop at the 6th (4 decreasing + 2 increasing calls) epoch
for epoch in 1:10
es() && break
end
```
The keyword argument `distance` of `early_stopping` is a function of the form `distance(best_score, score)`. By default `distance` is `-` and `init_score` is `Inf`, meaning that the monitored metric `f` is expected to be decreasing and minimized. If you use a metric such that improvement is shown by increasing values (e.g. accuracy), you can customize the `distance` function and the `init_score` value to, for example, `(best_score, score) -> score - best_score` and `-Inf`, respectively.
```julia
# create a pseudo-accuracy that increases by 0.01 each time from 0 to 1
# we call this like acc()
acc = let v = 0
() -> v = max(1, v + 0.01)
end
# create an early stopping trigger for accuracy
es = early_stopping(acc, 3; delta = (best_score, score) -> score - best_score, init_score = -Inf)
# this will iterate until the 10th epoch
for epoch in 1:10
es() && break
end
```
`early_stopping` and `plateau` are both built on top of `patience`. You can use `patience` to build your own triggers that use a patient counter. For example, if you want to trigger when the loss is below a threshold for several consecutive iterations:
```julia
threshold(f, thresh, delay) = patience(delay) do
f() < thresh
end
```
Both `predicate` in `patience` and `f` in `early_stopping` / `plateau` can accept extra arguments. You can pass such extra arguments to `predicate` or `f` through the returned function:
```julia
trigger = patience((a; b) -> a > b, 3)
# this will iterate until the 10th epoch
for epoch in 1:10
trigger(1; b = 2) && break
end
# this will stop at the 3rd epoch
for epoch in 1:10
trigger(3; b = 2) && break
end
```
```@docs
Flux.patience
Flux.early_stopping
Flux.plateau
```
---
### Src/Reference/Training/Gradients
```@meta
CollapsedDocStrings = true
```
# Automatic Differentiation in Flux
Flux's `gradient` function uses [Zygote](https://github.com/FluxML/Zygote.jl) by default, and also uses this function within [`train!`](@ref Flux.train!) to differentiate the model.
Zygote has its own [documentation](https://fluxml.ai/Zygote.jl/dev/), in particular listing some [important limitations](https://fluxml.ai/Zygote.jl/dev/limitations/).
Flux also has support for Enzyme.jl, documented [below](@ref autodiff-enzyme) and for Mooncake.jl.
## Generic Gradient Interface
```@docs
Flux.gradient(f, adtype::AbstractADType, args::Any...)
Flux.withgradient(f, adtype::AbstractADType, args::Any...)
```
## [Automatic Differentiation using Zygote.jl](@id autodiff-zygote)
The default AD backend in Flux is Zygote. Besides gradient calculation, Zygote also supports
higher-order derivatives, Jacobians, Hessians, and pullbacks.
```@docs
Zygote.jacobian(f, args...)
Zygote.withjacobian(f, args...)
Zygote.hessian
Zygote.hessian_reverse
Zygote.diaghessian
Zygote.pullback
```
## ChainRules for Zygote
Zygote uses [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl) to define how to differentiate functions.
Sometimes it is necessary to exclude some code, or a whole function, from automatic differentiation.
This can be done using the following methods:
```@docs
ChainRulesCore.ignore_derivatives
ChainRulesCore.@non_differentiable
```
To manually supply the gradient for one function, you should define a method of `rrule`. ChainRules has [detailed documentation](https://juliadiff.org/ChainRulesCore.jl/stable/) on how this works.
```@docs
ChainRulesCore.rrule
ChainRulesCore.frule
ChainRulesCore.@scalar_rule
ChainRulesCore.NoTangent
ChainRulesCore.ZeroTangent
ChainRulesCore.RuleConfig
ChainRulesCore.Tangent
ChainRulesCore.canonicalize
```
Gradient customization for other AD packages such as Enzyme and Mooncake has to be done according to their own documentation.
## [Automatic Differentiation using Enzyme.jl](@id autodiff-enzyme)
[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) is a new package for automatic differentiation.
Like Zygote.jl, calling `gradient(f, x)` causes it to hooks into the compiler and transform code that is executed while calculating `f(x)`, in order to produce code for `∂f/∂x`.
But it does so much later in the optimisation process (on LLVM instead of Julia's untyped IR) which you can [read about here](https://proceedings.nips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf)].
It needs far fewer custom rules than Zygote/ChainRules, and in particular is able to support mutation of arrays.
Flux now builds in support for this, using Enzyme's own `Duplicated` type.
Calling `Duplicated` on any Flux model which was defined using `@layer` will allocate space for the gradient,
and passing that to `gradient` (or `withgradient`, or `train!`) will then use Enzyme instead of Zygote.
The gradient functions still return the gradient as usual, which can then be passed to `update!`:
```julia-repl
julia> using Flux, Enzyme
julia> model = Chain(Dense(28^2 => 32, sigmoid), Dense(32 => 10), softmax); # from model zoo
julia> dup_model = Enzyme.Duplicated(model) # this allocates space for the gradient
Duplicated(
Chain(
Dense(784 => 32, σ), # 25_120 parameters
Dense(32 => 10), # 330 parameters
NNlib.softmax,
),
# norm(∇) ≈ 0.0f0
) # Total: 4 arrays, 25_450 parameters, 199.391 KiB.
julia> x1 = randn32(28*28, 1); # fake image
julia> y1 = [i==3 for i in 0:9]; # fake label
julia> grads_f = Flux.gradient((m,x,y) -> sum(abs2, m(x) .- y), dup_model, Const(x1), Const(y1)) # uses Enzyme
((layers = ((weight = Float32[-0.010354728 0.032972857 …
-0.0014538406], σ = nothing), nothing),), nothing, nothing)
```
The gradient returned here is also stored within `dup_model`.
Both share the same arrays -- what is returned is not a copy, just a view of the same memory (wrapped in `NamedTuple`s instead of `struct`s).
They will all be set to zero when you call `gradient` again, then replaced with the new values.
Alternatively, `gradient(f, args...; zero=false)` will add the new gradient to what's already stored.
Writing `Const(x1)` is optional, just plain `x1` is implicitly constant.
Any set of `Duplicated` and `Const` arguments may appear in any order, so long as there is at least one `Duplicated`.
The gradient `grads_f[1]` can be passed to `update!` as usual.
But for convenience, you may also use what is stored within `Duplicated`.
These are equivalent ways to perform an update step:
```julia-repl
julia> opt_state = Flux.setup(Adam(), model)
julia> ans == Flux.setup(Adam(), dup_model)
julia> Flux.update!(opt_state, model, grads_f[1]) # exactly as for Zygote gradients
julia> Flux.update!(opt_state, dup_model) # equivalent new path, Enzyme only
```
Instead of using these FLux functions, you can also use Enzyme's own functions directly.
`Enzyme.gradient` works like this:
```julia-repl
julia> grads_e = Enzyme.gradient(Reverse, (m,x,y) -> sum(abs2, m(x) .- y), model, Const(x1), Const(y1))
(Chain(Dense(784 => 32, σ), Dense(32 => 10), softmax), nothing, nothing)
julia> grads_f[1].layers[2].bias ≈ grads_e[1].layers[2].bias
true
```
Note that what `Enzyme.gradient` returns is an object like `deepcopy(model)` of the same type, `grads_e[1] isa Chain`.
But its fields contain the same gradient.
```@docs
Flux.gradient(f, args::Union{Flux.EnzymeCore.Const, Flux.EnzymeCore.Duplicated}...)
Flux.withgradient(f, args::Union{Flux.EnzymeCore.Const, Flux.EnzymeCore.Duplicated}...)
```
Enzyme.jl has [its own extensive documentation](https://enzymead.github.io/Enzyme.jl/stable/).
## Second-order AD
If you calculate a gradient within the loss function, then training will involve 2nd derivatives.
While this is in principle supported by Zygote.jl, there are many bugs, and Enzyme.jl is probably a better choice.
---
### Src/Reference/Training/Optimisers
```@meta
CurrentModule = Flux
CollapsedDocStrings = true
```
# [Optimisation Rules](@id man-optimisers)
Any optimization rule from Optimisers.jl can be used with [`train!`](@ref Flux.Train.train!) and
other training functions.
For full details of how the interface works, see the [Optimisers.jl documentation](https://fluxml.ai/Optimisers.jl/).
## Optimisers Reference
All optimisers return an object that, when passed to `train!`, will update the parameters passed to it.
```@docs
Optimisers.Descent
Optimisers.Momentum
Optimisers.Nesterov
Optimisers.RMSProp
Optimisers.Adam
Optimisers.RAdam
Optimisers.AdaMax
Optimisers.AdaGrad
Optimisers.AdaDelta
Optimisers.AMSGrad
Optimisers.NAdam
Optimisers.AdamW
Optimisers.OAdam
Optimisers.AdaBelief
Optimisers.Lion
```
## Composing Optimisers
Flux (through Optimisers.jl) defines a special kind of optimiser called `OptimiserChain` which takes in arbitrary optimisers as input. Its behaviour is similar to the usual optimisers, but differs in that it acts by calling the optimisers listed in it sequentially. Each optimiser produces a modified gradient
that will be fed into the next, and the resultant update will be applied to the parameter as usual. A classic use case is where adding decays is desirable. Optimisers.jl defines the basic decay corresponding to an $L_2$ regularization in the loss as `WeightDecay`.
```julia
opt = OptimiserChain(WeightDecay(1e-4), Descent())
```
Here we apply the weight decay to the `Descent` optimiser.
The resulting optimiser `opt` can be used as any optimiser.
```julia
w = [randn(10, 10), randn(10, 10)]
opt_state = Flux.setup(opt, w)
loss(w, x) = Flux.mse(w[1] * x, w[2] * x)
loss(w, rand(10)) # around 0.9
for t = 1:10^5
g = gradient(w -> loss(w[1], w[2], rand(10)), w)
Flux.update!(opt_state, w, g)
end
loss(w, rand(10)) # around 0.9
```
It is possible to compose optimisers for some added flexibility.
```@docs
Optimisers.OptimiserChain
```
## Decays
Similar to optimisers, Flux also defines some simple decays that can be used in conjunction with other optimisers, or standalone.
```@docs
Optimisers.SignDecay
Optimisers.WeightDecay
```
## Gradient Clipping
Gradient clipping is useful for training recurrent neural networks, which have a tendency to suffer from the exploding gradient problem. An example usage is
```julia
opt = OptimiserChain(ClipGrad(1e-3), Adam(1e-3))
```
```@docs
Optimisers.ClipGrad
Optimisers.ClipNorm
```
---
### Src/Reference/Training/Reference
```@meta
CollapsedDocStrings = true
```
# Training API Reference
The new version of Flux's training code was written as an independent package, [Optimisers.jl](https://github.com/FluxML/Optimisers.jl).
Only the function `train!` belongs to Flux itself.
The Optimisers package is designed to allow for immutable objects. But at present all Flux models contain parameter arrays (such as `Array`s and `CuArray`s) which can be updated in-place.
Because of this:
* The objects returned by `Optimisers.update!` can be ignored.
* Flux defines its own version of `setup` which checks this assumption.
(Using instead `Optimisers.setup` will also work, they return the same thing.)
The available optimization rules are listed the [optimisation rules](@ref man-optimisers) page here. See the [Optimisers documentation](https://fluxml.ai/Optimisers.jl/dev/) for details on how the rules work.
```@docs
Flux.Train.setup
Flux.Train.train!
Flux.Train.trainstep!
Flux.Train.trainstep_withgradient!
Optimisers.update
Optimisers.update!
Optimisers.setup
```
`train!` uses [`@progress`](https://github.com/JuliaLogging/ProgressLogging.jl) which should show a progress bar in VSCode automatically.
To see one in a terminal, you will need to install [TerminalLoggers.jl](https://github.com/JuliaLogging/TerminalLoggers.jl)
and follow its setup instructions.
There is also a method of `train!` which similarly takes `Duplicated(model)` and uses Enzyme.jl for differentiation (see (@ref autodiff-enzyme)):
```julia-repl
julia> opt_state = Flux.setup(Adam(0), model);
julia> Flux.train!((m,x,y) -> sum(abs2, m(x) .- y), dup_model, [(x1, y1)], opt_state)
```
## Optimisation Modifiers
The state returned by `setup` can be modified to temporarily prevent training of
some parts of the model, or to change the learning rate or other hyperparameter.
The functions for doing so may be accessed as `Flux.freeze!`, `Flux.thaw!`, and `Flux.adjust!`.
All mutate the state (or part of it) and return `nothing`.
```@docs
Optimisers.adjust!
Optimisers.freeze!
Optimisers.thaw!
```
---
### Src/Reference/Destructure
```@meta
CurrentModule = Flux
CollapsedDocStrings = true
```
# [Flat vs. Nested Structures](@id man-destructure)
A Flux model is a nested structure, with parameters stored within many layers. Sometimes you may want a flat representation of them, to interact with functions expecting just one vector. This is provided by `destructure`:
```julia-repl
julia> model = Chain(Dense(2=>1, tanh), Dense(1=>1))
Chain(
Dense(2 => 1, tanh), # 3 parameters
Dense(1 => 1), # 2 parameters
) # Total: 4 arrays, 5 parameters, 276 bytes.
julia> flat, rebuild = Flux.destructure(model)
(Float32[0.863101, 1.2454957, 0.0, -1.6345707, 0.0], Restructure(Chain, ..., 5))
julia> rebuild(zeros(5)) # same structure, new parameters
Chain(
Dense(2 => 1, tanh), # 3 parameters (all zero)
Dense(1 => 1), # 2 parameters (all zero)
) # Total: 4 arrays, 5 parameters, 276 bytes.
```
Both `destructure` and the `Restructure` function can be used within gradient computations. For instance, this computes the Hessian `∂²L/∂θᵢ∂θⱼ` of some loss function, with respect to all parameters of the Flux model. The resulting matrix has off-diagonal entries, which cannot really be expressed in a nested structure:
```julia-repl
julia> x = rand(Float32, 2, 16);
julia> grad = gradient(m -> sum(abs2, m(x)), model) # nested gradient
((layers = ((weight = Float32[10.339018 11.379145], bias = Float32[22.845667], σ = nothing), (weight = Float32[-29.565302;;], bias = Float32[-37.644184], σ = nothing)),),)
julia> function loss(v::Vector)
m = rebuild(v)
y = m(x)
sum(abs2, y)
end;
julia> gradient(loss, flat) # flat gradient, same numbers
(Float32[10.339018, 11.379145, 22.845667, -29.565302, -37.644184],)
julia> Zygote.hessian(loss, flat) # second derivative
5×5 Matrix{Float32}:
-7.13131 -5.54714 -11.1393 -12.6504 -8.13492
-5.54714 -7.11092 -11.0208 -13.9231 -9.36316
-11.1393 -11.0208 -13.7126 -27.9531 -22.741
-12.6504 -13.9231 -27.9531 18.0875 23.03
-8.13492 -9.36316 -22.741 23.03 32.0
julia> Flux.destructure(grad) # acts on non-models, too
(Float32[10.339018, 11.379145, 22.845667, -29.565302, -37.644184], Restructure(Tuple, ..., 5))
```
In order to collect all parameters of a model into a list instead, you can use the `trainables` function:
```julia-repl
julia> Flux.trainables(model)
5-element Vector{AbstractArray}:
[0.863101 1.2454957]
[0.0]
[1.290355429422727;;]
[0.0]
```
Any mutation of the elements of the resulting list will affect the model's parameters.
## All Parameters
The functions `destructure` and `trainables` live in [`Optimisers.jl`](https://github.com/FluxML/Optimisers.jl).
```@docs
Optimisers.destructure
Optimisers.trainable
Optimisers.trainables
Optimisers.isnumeric
Flux.params
```
## All Layers
Another kind of flat view of a nested model is provided by the `modules` command. This extracts a list of all layers:
```@docs
Flux.modules
```
## Save and Load
```@docs
Flux.state
Flux.loadmodel!
```
## KeyPath
```@docs
Functors.KeyPath
Functors.getkeypath
Functors.haskeypath
Functors.setkeypath!
```
---
### Src/Reference/Outputsize
```@meta
CollapsedDocStrings = true
```
# Shape Inference
Flux has some tools to help generate models in an automated fashion, by inferring the size
of arrays that layers will receive, without doing any computation.
This is especially useful for convolutional models, where the same [`Conv`](@ref) layer
accepts any size of image, but the next layer may not.
The higher-level tool is a macro [`@autosize`](@ref) which acts on the code defining the layers,
and replaces each appearance of `_` with the relevant size. This simple example returns a model
with `Dense(845 => 10)` as the last layer:
```julia
@autosize (28, 28, 1, 32) Chain(Conv((3, 3), _ => 5, relu, stride=2), Flux.flatten, Dense(_ => 10))
```
The input size may be provided at runtime, like `@autosize (sz..., 1, 32) Chain(Conv(`..., but all the
layer constructors containing `_` must be explicitly written out -- the macro sees the code as written.
This macro relies on a lower-level function [`outputsize`](@ref Flux.outputsize), which you can also use directly:
```julia
c = Conv((3, 3), 1 => 5, relu, stride=2)
Flux.outputsize(c, (28, 28, 1, 32)) # returns (13, 13, 5, 32)
```
The function `outputsize` works by passing a "dummy" array into the model, which propagates through very cheaply.
It should work for all layers, including custom layers, out of the box.
An example of how to automate model building is this:
```jldoctest; output = false, setup = :(using Flux)
"""
make_model(width, height, [inchannels, nclasses; layer_config])
Create a CNN for a given set of configuration parameters. Arguments:
- `width`, `height`: the input image size in pixels
- `inchannels`: the number of channels in the input image, default `1`
- `nclasses`: the number of output classes, default `10`
- Keyword `layer_config`: a vector of the number of channels per layer, default `[16, 16, 32, 64]`
"""
function make_model(width, height, inchannels = 1, nclasses = 10;
layer_config = [16, 16, 32, 64])
# construct a vector of layers:
conv_layers = []
push!(conv_layers, Conv((5, 5), inchannels => layer_config[1], relu, pad=:same))
for (inch, outch) in zip(layer_config, layer_config[2:end])
push!(conv_layers, Conv((3, 3), inch => outch, sigmoid, stride=2))
end
# compute the output dimensions after these conv layers:
conv_outsize = Flux.outputsize(conv_layers, (width, height, inchannels); padbatch=true)
# use this to define appropriate Dense layer:
last_layer = Dense(prod(conv_outsize) => nclasses)
return Chain(conv_layers..., Flux.flatten, last_layer)
end
m = make_model(28, 28, 3, layer_config = [9, 17, 33, 65])
Flux.outputsize(m, (28, 28, 3, 42)) == (10, 42) == size(m(randn(Float32, 28, 28, 3, 42)))
# output
true
```
Alternatively, using the macro, the definition of `make_model` could end with:
```
# compute the output dimensions & construct appropriate Dense layer:
return @autosize (width, height, inchannels, 1) Chain(conv_layers..., Flux.flatten, Dense(_ => nclasses))
end
```
### Listing
```@docs
Flux.@autosize
Flux.outputsize
```
---
### Src/Reference/Utilities
```@meta
CurrentModule = Flux
CollapsedDocStrings = true
```
# [Random Weight Initialisation](@id man-init-funcs)
Flux initialises convolutional layers and recurrent cells with `glorot_uniform` by default.
Most layers accept a function as an `init` keyword, which replaces this default. For example:
```jldoctest; setup = :(using Flux)
julia> conv = Conv((3, 3), 3 => 2, relu; init=Flux.glorot_normal)
Conv((3, 3), 3 => 2, relu) # 56 parameters
julia> conv.bias
2-element Vector{Float32}:
0.0
0.0
```
Note that `init` creates the weight array, but not the bias vector.
Many of the initialisation functions accept keywords such as `gain`,
and a random number generator. To make it easy to pass these to layers,
there are methods which return a function:
```jldoctest; setup = :(using Flux, Random)
julia> Dense(4 => 5, tanh; init=Flux.glorot_uniform(gain=2))
Dense(4 => 5, tanh) # 25 parameters
julia> Dense(4 => 5, tanh; init=Flux.randn32(MersenneTwister(1)))
Dense(4 => 5, tanh) # 25 parameters
```
## Initialisation functions
```@docs
Flux.glorot_uniform
Flux.glorot_normal
Flux.kaiming_uniform
Flux.kaiming_normal
Flux.truncated_normal
Flux.lecun_normal
Flux.orthogonal
Flux.sparse_init
Flux.identity_init
Flux.ones32
Flux.zeros32
Flux.rand32
Flux.randn32
Flux.create_bias
```
These functions call:
```@docs
Flux.rng_from_array
Flux.nfan
```
## Changing the type of all parameters
The default `eltype` for models is `Float32` since models are often trained/run on GPUs.
The `eltype` of model `m` can be changed to `Float64` by `f64(m)`:
```@docs
Flux.f64
Flux.f32
Flux.f16
Flux.bf16
```
---
### Src/Tutorials/Custom Layers
# [Defining Customised Layers](@id man-advanced)
Here we will try and describe usage of some more advanced features that Flux provides to give more control over model building.
## Custom Model Example
Here is a basic example of a custom model. It simply adds the input to the result from the neural network.
```julia
struct CustomModel{T <: Chain} # Parameter to avoid type instability
chain::T
end
function (m::CustomModel)(x)
# Arbitrary code can go here, but note that everything will be differentiated.
# Zygote does not allow some operations, like mutating arrays.
return m.chain(x) + x
end
# This is optional but recommended for pretty printing and other niceties
Flux.@layer CustomModel
```
Notice that we parameterized the type of the `chain` field. This is necessary for fast Julia code, so that that struct field can be given a concrete type. `Chain`s have a type parameter fully specifying the types of the layers they contain. By using a type parameter, we are freeing Julia to determine the correct concrete type, so that we do not need to specify the full, possibly quite long, type ourselves.
You can then use the model like:
```julia
chain = Chain(Dense(10 => 10, relu), Dense(10 => 10))
model = CustomModel(chain)
model(rand(Float32, 10))
```
For an intro to Flux and automatic differentiation, see this [tutorial](https://fluxml.ai/tutorials/2020/09/15/deep-learning-flux.html).
## Customising Parameter Collection for a Model
Taking reference from our example `Affine` layer from the [basics](@ref man-basics).
By default all the fields in the `Affine` type are collected as its parameters, however, in some cases it may be desired to hold other metadata in our "layers" that may not be needed for training, and are hence supposed to be ignored while the parameters are collected. With Flux, the way to mark some fields of our layer as trainable is through overloading the `trainable` function:
```julia-repl
julia> struct Affine
W
b
end
julia> Affine(in::Int, out::Int) = Affine(randn(out, in), randn(out));
julia> (m::Affine)(x) = m.W * x .+ m.b;
julia> Flux.@layer Affine
julia> a = Affine(Float32[1 2; 3 4; 5 6], Float32[7, 8, 9])
Affine(Float32[1.0 2.0; 3.0 4.0; 5.0 6.0], Float32[7.0, 8.0, 9.0])
julia> Flux.trainable(a) # default behavior
(W = Float32[1.0 2.0; 3.0 4.0; 5.0 6.0], b = Float32[7.0, 8.0, 9.0])
julia> Flux.trainable(a::Affine) = (; W = a.W) # returns a NamedTuple using the field's name
julia> Flux.trainable(a)
(W = Float32[1.0 2.0; 3.0 4.0; 5.0 6.0],)
```
Only the fields returned by `trainable` will be seen by `Flux.setup` and `Flux.update!` for training. But all fields will be seen by `gpu` and similar functions, for example:
```julia-repl
julia> a |> f16
Affine(Float16[1.0 2.0; 3.0 4.0; 5.0 6.0], Float16[7.0, 8.0, 9.0])
```
Note that there is no need to overload `trainable` to hide fields which do not contain numerical array (for example, activation functions, or Boolean flags). These are always ignored by training.
The exact same method of `trainable` can also be defined using the macro, for convenience:
```julia
Flux.@layer Affine trainable=(W,)
```
There is a second, more severe, kind of restriction possible. This is not recommended, but is included here for completeness. Calling `Functors.@functor Affine (W,)` means that no exploration of the model will ever visit the other fields: They will not be moved to the GPU by [`gpu`](@ref), and their precision will not be changed by `f32`. This requires the `struct` to have a corresponding constructor that accepts only `W` as an argument.
## Custom multiple input or output layer
Sometimes a model needs to receive several separate inputs at once or produce several separate outputs at once. In other words, there multiple paths within this high-level layer, each processing a different input or producing a different output. A simple example of this in machine learning literature is the [inception module](https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Szegedy_Rethinking_the_Inception_CVPR_2016_paper.pdf).
We could have a struct that stores the weights of along each path and implement the joining/splitting in the forward pass function. That would mean a new struct for each different block,
e.g. one would have a `TransformerBlock` struct for a transformer block, and a `ResNetBlock` struct for a ResNet block, each block being composed by smaller sub-blocks. This is often the simplest and cleanest way to implement complex models.
This guide instead will show you how to construct a high-level layer (like [`Chain`](@ref)) that is made of multiple sub-layers for each path. It may be the case that using the layers described as follows makes the definition of your model harder to read and to change. In that case, consider using the simpler approach of defining a custom structure described above.
### Multiple inputs: a custom `Join` layer
Our custom `Join` layer will accept multiple inputs at once, pass each input through a separate path, then combine the results together. Note that this layer can already be constructed using [`Parallel`](@ref), but we will first walk through how do this manually.
We start by defining a new struct, `Join`, that stores the different paths and a combine operation as its fields.
```julia
using Flux
using CUDA
# custom join layer
struct Join{T, F}
combine::F
paths::T
end
# allow Join(op, m1, m2, ...) as a constructor
Join(combine, paths...) = Join(combine, paths)
```
Notice again that we parameterized the type of the `combine` and `paths` fields. In addition to the performance considerations of concrete types, this allows either field to be `Vector`s, `Tuple`s, or one of each - we don't need to pay attention to which.
The next step is to use [`Flux.@layer`](@ref) to make our struct behave like a Flux layer.
In Flux < v0.15 this used to be important so that calling `Flux.setup` on a `Join` maps over the underlying trainable arrays on each path. Since Flux v0.15, this is no longer necessary, since now Functors.jl automatically traverses custom types. However, [`Flux.@layer`](@ref) is still recommended for pretty printing and other niceties.
```julia
Flux.@layer Join
```
Finally, we define the forward pass. For `Join`, this means applying each `path` in `paths` to each input array, then using `combine` to merge the results.
```julia
(m::Join)(xs::Tuple) = m.combine(map((f, x) -> f(x), m.paths, xs)...)
(m::Join)(xs...) = m(xs)
```
Lastly, we can test our new layer. Thanks to the proper abstractions in Julia, our layer works on GPU arrays out of the box!
```julia
model = Chain(
Join(vcat,
Chain(Dense(1 => 5, relu), Dense(5 => 1)), # branch 1
Dense(1 => 2), # branch 2
Dense(1 => 1) # branch 3
),
Dense(4 => 1)
) |> gpu
xs = map(gpu, (rand(1), rand(1), rand(1)))
model(xs)
# returns a single float vector with one value
```
!!! note
This `Join` layer is available from the [Fluxperimental.jl](https://github.com/FluxML/Fluxperimental.jl) package.
#### Using `Parallel`
Flux already provides [`Parallel`](@ref) that can offer the same functionality. In this case, `Join` is going to just be syntactic sugar for `Parallel`.
```julia
Join(combine, paths) = Parallel(combine, paths)
Join(combine, paths...) = Join(combine, paths)
# use vararg/tuple version of Parallel forward pass
model = Chain(
Join(vcat,
Chain(Dense(1 => 5, relu), Dense(5 => 1)),
Dense(1 => 2),
Dense(1 => 1)
),
Dense(4 => 1)
) |> gpu
xs = map(gpu, (rand(1), rand(1), rand(1)))
model(xs)
# returns a single float vector with one value
```
### Multiple outputs: a custom `Split` layer
Our custom `Split` layer will accept a single input, then pass the input through a separate path to produce multiple outputs.
We start by following the same steps as the `Join` layer: define a struct, use [`Flux.@layer`](@ref), and define the forward pass.
```julia
using Flux
using CUDA
# custom split layer
struct Split{T}
paths::T
end
Split(paths...) = Split(paths)
Flux.@layer Split
(m::Split)(x::AbstractArray) = map(f -> f(x), m.paths)
```
Now we can test to see that our `Split` does indeed produce multiple outputs.
```julia
model = Chain(
Dense(10 => 5),
Split(Dense(5 => 1, tanh), Dense(5 => 3, tanh), Dense(5 => 2))
) |> gpu
model(gpu(rand(10)))
# returns a tuple with three float vectors
```
A custom loss function for the multiple outputs may look like this:
```julia
using Statistics
# assuming model returns the output of a Split
# x is a single input
# ys is a tuple of outputs
function loss(x, ys, model)
# rms over all the mse
ŷs = model(x)
return sqrt(mean(Flux.mse(y, ŷ) for (y, ŷ) in zip(ys, ŷs)))
end
```
!!! note
This `Split` layer is available from the [Fluxperimental.jl](https://github.com/FluxML/Fluxperimental.jl) package.
---
Flux's layers