Training an NST model

In a nutshell, ML models are similar to functions : they take one or multiple inputs to returns one or multiple outputs. One big difference though is that as a developer you don’t write a model, you train it. The trained model can then be fed inputs which will be processed through its underlying tensors to reach a result, or prediction.

Our knowledge is therefore shifted to a higher level task : implementing the training algorithm.

Multiple frameworks are available to achieve this : Tensorflow, Keras, Caffe, PyTorch… We will focus on that last one in this article, PyTorch, as it provides a strong GPU acceleration, which will reveal as an important feature in the next steps.

To get started, I recommend using Anaconda which makes it very easy to manage your environment in case you have several conflicting dependencies on your workstation. If you have an NVidia Graphic Card, it’s very important that you also install the proper CUDA toolkit for it to enable GPU acceleration during training.

Once your environment is set, you can go to PyTorch Github repo which list multiple usage examples, one being a Fast Neural Style sample. Clone this repo onto your workstation and activate your environment.

You should now be able to try training your first model using the default parameters. This will start a training loop which will save a checkpoint every two thousands iterations.

Training a model with PyTorch

Those checkpoint are nothing more than a model at a given number of iterations. Usually the more iteration, the better, but in our case we are aiming for beauty, which is subjective and difficult to measure without actually having a look at the result.

You can look at how neural_style.py implements the actual training and what makes it “improve” over the previous iteration, which uses multiple metrics like style-weight, content-weight, regression loss and more.

Stylizing an image with PyTorch

Trying with other style image, you‘ll find that tweaking those style-weight and content-weight is a good place to start in order to improve the model predictions.

Exporting your model to CoreML

Here comes one of the tricky parts. PyTorch doesn’t natively support exports to CoreML (yet). Fortunately, you might have heard about an endeavour aiming at make Machine Learning models interoperable, the Open Neural Network Exchange.

ONNX provides tools for importing and exporting models from almost any framework to the .onnx format (mandatory xkcd).

Hence, we can export our model to CoreML in 2 phases :

PyTorch → ONNX→ CoreML

neural_style.py already have an ONNX export, so we really just need to implement the second step. We can find everything we need on the onnx-coreml Github repo to bridge that gap.

ONNX to CoreML export

Now that we have our ONNX to CoreML converter, we can convert a checkpoint with a simple bash command.

Performing a PyTorch → ONNX→ CoreML export

Using a CoreML NST model on iOS

You should feel more at home for this part if, like me, you are an iOS developer.

First, we start by creating a simple project with an UIImageView to hold our original and output image, as well as an UIButton to launch the process.

Next, we add an UIBarButtonItem to pick an image from the device’s photo library and set it onto the UIImageView.

Alternatively, we can just provide a default image and use that one.