# Simple ConvNet

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:

using Flux, MLDatasets, Statistics
using Flux: onehotbatch, onecold, logitcrossentropy
using MLDatasets: MNIST
using Base.Iterators: partition
using Printf, BSON
using Parameters: @with_kw
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:

@with_kw mutable struct Args
lr::Float64 = 3e-3
epochs::Int = 20
batch_size = 128
savepath::String = "./"
end


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):

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.

function get_processed_data(args)
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


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.

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
MaxPool((2,2)),

# Second convolution, operating upon a 14x14 image
MaxPool((2,2)),

# Third convolution, operating upon a 7x7 image
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. It enables us to call the layers in sequence on a given input. Also, we use the function flatten to reshape the output image from the last convolution layer. Finally, we call the Dense function to create the classification layer.

Before training our model, we need to define a few functions that will be helpful for the process:

• augment augments the data by adding gaussian random noise to our image to make it more robust:

  augment(x) = x .+ gpu(0.1f0*randn(eltype(x), size(x)))


• anynan checks whether any element of the params is NaN or not:

   anynan(x) = any(y -> any(isnan, y), x)


• accuracy computes the accuracy of our ConvNet:

   accuracy(x, y, model) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))


Finally, we define the train function:

function train(; kws...)
args = Args(; kws...)

train_set, test_set = get_processed_data(args)

# Define our model.  We will use a simple convolutional architecture with
# three iterations of Conv -> ReLU -> MaxPool, followed by a final Dense layer.
@info("Building model...")
model = build_model(args)

# Load model and datasets onto GPU, if enabled
train_set = gpu.(train_set)
test_set = gpu.(test_set)
model = gpu(model)

# Make sure our model is nicely precompiled before starting our training loop
model(train_set[1][1])

# loss() calculates the crossentropy loss between our prediction y_hat
# (calculated from model(x)) and the ground truth y.  We augment the data
# a bit, adding gaussian random noise to our image to make it more robust.
function loss(x, y)
x̂ = augment(x)
ŷ = model(x̂)
return logitcrossentropy(ŷ, y)
end

# Train our model with the given training set using the ADAM optimizer and
# printing out performance against the test set as we go.

@info("Beginning training loop...")
best_acc = 0.0
last_improvement = 0
for epoch_idx in 1:args.epochs
# Train for a single epoch
Flux.train!(loss, params(model), train_set, opt)

# Terminate on NaN
if anynan(Flux.params(model))
@error "NaN params"
break
end

# Calculate accuracy:
acc = accuracy(test_set..., model)

@info(@sprintf("[%d]: Test accuracy: %.4f", epoch_idx, acc))
# If our accuracy is good enough, quit out.
if acc >= 0.999
@info(" -> Early-exiting: We reached our target accuracy of 99.9%")
break
end

# If this is the best accuracy we've seen so far, save the model out
if acc >= best_acc
@info(" -> New best accuracy! Saving model out to mnist_conv.bson")
BSON.@save joinpath(args.savepath, "mnist_conv.bson") params=cpu.(params(model)) epoch_idx acc
best_acc = acc
last_improvement = epoch_idx
end

# If we haven't seen improvement in 5 epochs, drop our learning rate:
if epoch_idx - last_improvement >= 5 && opt.eta > 1e-6
opt.eta /= 10.0
@warn(" -> Haven't improved in a while, dropping learning rate to \$(opt.eta)!")

# After dropping learning rate, give it a few epochs to improve
last_improvement = epoch_idx
end

if epoch_idx - last_improvement >= 10
@warn(" -> We're calling this converged.")
break
end
end
end


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:

• 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 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.

Finally, to test our model we define the test function:

function test(; kws...)
args = Args(; kws...)

_,test_set = get_processed_data(args)

# Re-constructing the model with random initial weights
model = build_model(args)


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%).