Secret recipes for turi create style transfer model
I had the opportunity to train a custom style transfer model for an AR app called notable me that would swap faces into bills.
Notable Me App on App Store: https://apps.apple.com/app/notable-me/id1453810473
I want to share the process that went through.
Target effect
The notable me app swaps the face of a selfie with the portrait of a real bill.
We wanted to make the app more magical by making the textures of the selfie image similar to the bills face texture style.
Why turi create?
Turi Create is Apple’s open source project that simplifies the development of custom machine learning models.
I tried using some models in github written with pytorch, but when I converted to trained model to CoreML and uploaded it to the iPhone, the app crashed because of excessive memory usage, so I decided to use the style transfer model of turi create that Apple provides for app developers who are non machine learning expert.
Prepare environment
- style images: Images of banknote
- content images: selfie dataset (46,836 images, 1GB)
- https://www.crcv.ucf.edu/data/Selfie/
Find and Use right variables
If you train the model with default setting, the result would not great. There is more parameters you can play with than official documentation. Here is parameters what I found that you should know for better results.
Use right amount of GPU
If one or more GPUs are available on your machine, set it to use two of them. I tried some experimentation and results show that using two GPUs is the best option.
You can set it as follows.
turicreate.config.set_num_gpus(2)batch_size
The model will use the number of images specified in the batch size in one learning step. Although it is generally said that using the largest batch size as the graphics card’s GPU memory permits speeds up learning, there is also the opinion that setting a small batch size depending on the situation has a good effect on the performance of the final model.
So I decided to choose a batch size that seems to reduce the loss value the fastest by a simple experiment.
Thus, I used `batch_size` as 6. You can set it like below.
model = turicreate.style_transfer.create(
styles, content,
batch_size=6,
)pretrained_weights
Default value is False. If set to True, training will start from the pre-trained transfomer model. or transfomer model will learn from the scratch.
- pros: The loss will reduce rapidly. This is big plus.
- cons: If your style image is experimental, pre-trained transfomer model will not learn your unique style well.
Try True first, try False if it does not work.
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
pretrained_weights: True # <= try this first
}
)max_iterations
The bigger the better. However, you should train at least until the loss value is no longer reduced. Even after the loss value is stagnant, the quality may be improved if you continue learning.
model = turicreate.style_transfer.create(
styles, content,
max_iterations=50000,
)print_loss_breakdown
With the `print_loss_breakdown` option, style loss and content loss will be shown. The loss values can be obtained by evaluating the generated image.
- style loss: How well did the model reflect texture of the style image?
- content loss: How well did the model preserved the shapes of the content image?
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
print_loss_breakdown: True
}
)finetune_all_params
When set True, not only transformer model but also feature extractor model will be trained. If the content dataset is small, the fine-tuned model may show unstable results depending on the new input image.
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
finetune_all_params: True # set this if you have plenty of content images
}
)use_augmentation
This option will augment content images by resizing croping rotating etc.
The richer the content images, the more stable the model will be. I strongly recommend using this option.
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
use_augmentation: True # always use this.
}
)
style_loss_mult and content_loss_mult
Since the values of `style_loss_mult` and` content_loss_mult` affect each other in model learning, let’s fix the value of `content_loss_mult` to 1 and explore the value of `style_loss_mult` to a value between 0 and 1.
If the value of style_loss_mult is high, the shape of the content is not preserved. If the value is low, the style does not show. It is important to find a combination of values that show the exact effect I want.
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
style_loss_mult: [1e-2, 1e-3, 1e-3, 1e-4], # What I found for my model
}
)input_shape
The content and style images are resized to the size specified in `input_shape` before being used in the model.
So, if you do not see the texture details with resized image, you need to crop the style image or increase the size of the input_shape. The default value is (256, 256).
model = turicreate.style_transfer.create(
styles, content,
_advanced_parameters= {
input_shape: (488, 488), # ← What I found for my model
}
)Style Image Crop
You need to modify the style image so that unwanted colors or textures are not reflected in the generated image. In this case, the model generated unwanted yellowish color in the face.
Training code with advanced parameters
Result
