NOTE: This version works for Julia 0.6. An update to support Julia 1.0 and latest Flux, is WIP. Check the julia-1.0 branch for latest updates. It contains working code which needs to be trained. The models for the last release will not work in julia-1.0 branch. But soon new and better pre-trained models will be added.
Implementation of Fast Neural Style Transfer in Julia using Flux.jl
To install this package simply run
Pkg.clone("https://github.com/avik-pal/FastStyleTransfer.jl")
Using it is very simple. You just need to worry about 2 functions.
The first function is train
. You need this in case you are training on a new style. If you ever end up training on a new style please send in a PR with the model. Now the arguments you need to pass to the function are:
train_data_path
: Path to the MS COCO Dataset.batch_size
η
: Learning rate. Ideally set it to0.001
.style_image_path
epochs
model_save_path
content_weight
: The priority you want to assign to the content. A higher value leads to a better retainment of the original features.style_weight
: Same as content weight only for the style image.model
: PassTransformerNet
orTransformerNet2
. You can pass a custom Network as well.images
: Total images from the COCO dataset that you want to train you model on.
The other function would be stylize
. This will probably will be the only function most people shall use. So lets see its arguments:
image_path
model_path
save_path
: Can be left blank. The image will be stored in the same directory as the image with a stylized tag.display_img
: If running from the terminal set it tofalse
. Comes handy when you want to see the image without having to open thesave_path
.
Some speed statistics: The model runs quite fast taking only 5ms
to stylize a 640 x 640
image on a P100 GPU.
- Implement a GPU Kernel for Instance Normalization
- Substitute Zero Padding for Reflection Padding
- Implement the GPU Kernel for Upsampling Layer
- Perceptual Losses for Real-Time Style Transfer and Super-Resolution
- A Neural Algorithm of Artistic Style
- Instance Normalization: The Missing Ingredient for Fast Stylization
If you are interested in the implementation look up this blog post.