Deep Learning with Julia, Flux.jl story | by Artur Jurgas | Towards Data Science

A new challenger in the Data Science scene appears: Julia. It’s fast, it’s easy to type, it has a great documentation and community.
But it lacks examples and tutorials to learn from so in this article we will build a classic: MNIST classifier using Convolutional Neural Network.
I will focus on more beginner friendly style for those of you who know just a bit or two about Deep Learning but are curious to try a new language. I will try to explain Julia syntax where it differs from Python.
Flux.jl. It is one of the main deep learning libraries for Julia. You can view some examples here although they can be somewhat intimidating for a beginner.
While PyTorch or TensorFlow already provides a great ecosystem for Deep Learning in Python they are mostly written in C++ or even Cuda for the sweet GPU performance.
So if you want to write some custom code to experiment a bit It can get quite complicated (although it’s still probably a better idea to write production code in them).
Julia on the other hand provides you with speed from the start. So if you want to write some custom loss function you can do it in Julia. Combine it with little effort to put it on the GPU and clear source code and you have something quite appealing even for a beginner.
Importing packages is pretty straightforward. We will import Flux (of course), Statistics and MLDatasets for a simpler data preparation.
using Fluxusing Flux: Data.DataLoaderusing Flux: onehotbatch, onecold, crossentropyusing Flux: @epochsusing Statisticsusing MLDatasets# Load the datax_train, y_train = MLDatasets.MNIST.traindata()x_valid, y_valid = MLDatasets.MNIST.testdata()# Add the channel layerx_train = Flux.unsqueeze(x_train, 3)x_valid = Flux.unsqueeze(x_valid, 3)# Encode labelsy_train = onehotbatch(y_train, 0:9)y_valid = onehotbatch(y_valid, 0:9)# Create the full datasettrain_data = DataLoader(x_train, y_train, batchsize=128)
(Pro tip: Julia will by default print the output of the function. You can suppress that but typing “;” at the end)
Flux will expect our image data to be in WHCN order (width, height, # channels, batch size) so we will have to add a channel layer. Thankfully, there is already a function for that called unsqueeze.
Later we will use cross entropy loss, so we need to also encode our labels here using onehotbatch.
Our model will have 8 layers. 4 of them will be convolutions with Relu and then we will do mean pooling, flatten it and finally feed it into a linear layer with softmax at the end.
We can ‘chain’ everything together using the Chain function
model = Chain(    # 28x28 => 14x14    Conv((5, 5), 1=>8, pad=2, stride=2, relu),    # 14x14 => 7x7    Conv((3, 3), 8=>16, pad=1, stride=2, relu),    # 7x7 => 4x4    Conv((3, 3), 16=>32, pad=1, stride=2, relu),    # 4x4 => 2x2    Conv((3, 3), 32=>32, pad=1, stride=2, relu),        # Average pooling on each width x height feature map    GlobalMeanPool(),    flatten,        Dense(32, 10),    softmax)
Every convolution layer takes an image and creates channel layers out of it. So, the next convolution layer will take a smaller image with more channels in it. We also apply padding on the first layer and a stride of 2 on all of them.
If you need a refresher look at the gif below and this article. Padding basically makes our image bigger for the kernel and stride defines how big step it takes.
Then there is the mean pooling layer. It takes the feature map from the convolutional layers and takes a mean from every channel.
But before we can feed the data further into the linear layer, we must get rid of the 1x1 dimensions called singletons. That is what the flatten layer is for.
If you want to see how the dimensions change on every layer I will provide a link to a jupyter notebook at the end of the article.
Now is the time where we can already feed our data to the model and get predictions. They will not be any useful yet, but it is a good way to check if we have done everything correctly. To decode the prediction, use onecold function.
# Getting predictionsŷ = model(x_train)# Decoding predictionsŷ = onecold(ŷ)println("Prediction of first image: $(ŷ[1])")
Now is time to choose how we will update the models parameters and how we will check its performance.
accuracy(ŷ, y) = mean(onecold(ŷ) .== onecold(y))loss(x, y) = Flux.crossentropy(model(x), y)# learning ratelr = 0.1opt = Descent(lr)ps = Flux.params(model)
We have our standard accuracy metric and cross entropy loss from Flux. For an optimizer we will choose a simple gradient descent with a learning rate of 0.1. There are of course much better choices as Momentum or ADAM, but it will suffice for a simple guide.
The way Flux’s differentiation library Zygote works is somewhat different than the ones used in PyTorch for example. It is worth looking into on its own but for now we only need to know that we have to get the parameters out of our model.
To do it we simply call the params function with our model as an input.
And now the thing we were waiting for: training the model.
number_epochs = 10@epochs number_epochs Flux.train!(loss, ps, train_data, opt)accuracy(model(x_train), y_train)
That's it. That’s the training. We call the train! function with the loss, parameters, data, and optimizer. We use the @epochs macro with the number of how many times we want it to execute (by default it only does one).
The “!” in the function name usually means that the function creates side-efects or that it does things in-place (if you use numpy in Python).
“@” precedes macros. The are a gateway to what is called meta programming. They change the code of the function so you can kind of compare it to decorators in Pyton (though they are not the same thing strictly speaking).
Now let’s write the training loop ourselves. And the best part is that if we were to compare the time it takes to execute both it would be actually similar.
From the documentation:
The Flux.train! function can be very convenient, especially for simple problems. Its also very flexible with the use of callbacks. But for some problems its much cleaner to write your own custom training loop.
So, let's do just that.
for batch in train_data        gradient = Flux.gradient(ps) do      # Remember that inside the loss() is the model      # `...` syntax is for unpacking data      training_loss = loss(batch...)      return training_loss    end        Flux.update!(opt, ps, gradient)end
We loop over the training dataset (from DataLoader). Then we use the do keyword to map the loss calculations to the gradient function. We then call the update! with the optimizer, parameters and saved gradient and it is done.
Oh, you don’t like the update! function? No problem. Let us write our own loop.
for x in ps    x .-= lr .* gradient[x] # Update parametersend
The . before each sign is telling Julia to do the operation element-wise.
If you look into the source code of Flux you will see exactly that, just in two functions. Don’t believe me? See for yourself here and here.
And that is what is most appealing in Julia. It will be quick, and it is ready for the GPU. Awesome!
Want more? Let’s create some callbacks. They are essentially functions that will be called while training.
loss_vector = Vector{Float64}()callback() = push!(loss_vector, loss(x_train, y_train))Flux.train!(loss, ps, train_data, opt, cb=callback)
You can treat push! as an append function from numpy.
And now we have a list with the loss after each batch of the data (keep in mind that this is a lot of calculations for bigger datasets).
Now you know how to use Flux in Julia. You even know how to write efficient custom functions for your model. And how long has it taken? Not much, hasn’t it?
Make no mistake, Julia's and Flux’s days are still pretty early but if you like the language and how the packages are written I say it is worth to experiment with them a bit.
After all we all write better code in our favorite language if we have some experience with others.
You can see a jupyter notebook with all the code in it here.
Come hang out with me on twitter and thank you for reading!