CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/generative/adain/" />  <meta property="og:title" content="Keras documentation: Neural Style Transfer with AdaIN"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Neural Style Transfer with AdaIN"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Neural Style Transfer with AdaIN</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink active" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink2" href="/examples/generative/ddim/">Denoising Diffusion Implicit Models</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion/">A walk through latent space with Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/dreambooth/">DreamBooth</a> <a class="nav-sublink2" href="/examples/generative/ddpm/">Denoising Diffusion Probabilistic Models</a> <a class="nav-sublink2" href="/examples/generative/fine_tune_via_textual_inversion/">Teach StableDiffusion new concepts via Textual Inversion</a> <a class="nav-sublink2" href="/examples/generative/finetune_stable_diffusion/">Fine-tuning Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/vae/">Variational AutoEncoder</a> <a class="nav-sublink2" href="/examples/generative/dcgan_overriding_train_step/">GAN overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/wgan_gp/">WGAN-GP overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/conditional_gan/">Conditional GAN</a> <a class="nav-sublink2" href="/examples/generative/cyclegan/">CycleGAN</a> <a class="nav-sublink2" href="/examples/generative/gan_ada/">Data-efficient GANs with Adaptive Discriminator Augmentation</a> <a class="nav-sublink2" href="/examples/generative/deep_dream/">Deep Dream</a> <a class="nav-sublink2" href="/examples/generative/gaugan/">GauGAN for conditional image generation</a> <a class="nav-sublink2" href="/examples/generative/pixelcnn/">PixelCNN</a> <a class="nav-sublink2" href="/examples/generative/stylegan/">Face image generation with StyleGAN</a> <a class="nav-sublink2" href="/examples/generative/vq_vae/">Vector-Quantized Variational Autoencoders</a> <a class="nav-sublink2" href="/examples/generative/neural_style_transfer/">Neural style transfer</a> <a class="nav-sublink2 active" href="/examples/generative/adain/">Neural Style Transfer with AdaIN</a> <a class="nav-sublink2" href="/examples/generative/gpt2_text_generation_with_keras_hub/">GPT2 Text Generation with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_gpt/">GPT text generation from scratch with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_with_miniature_gpt/">Text generation with a miniature GPT</a> <a class="nav-sublink2" href="/examples/generative/lstm_character_level_text_generation/">Character-level text generation with LSTM</a> <a class="nav-sublink2" href="/examples/generative/text_generation_fnet/">Text Generation using FNet</a> <a class="nav-sublink2" href="/examples/generative/molecule_generation/">Drug Molecule Generation with VAE</a> <a class="nav-sublink2" href="/examples/generative/wgan-graphs/">WGAN-GP with R-GCN for the generation of small molecular graphs</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion_3/">A walk through latent space with Stable Diffusion 3</a> <a class="nav-sublink2" href="/examples/generative/real_nvp/">Density estimation using Real NVP</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparameter Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> <a class="nav-link" href="/keras_cv/" role="tab" aria-selected="">KerasCV: Computer Vision Workflows</a> <a class="nav-link" href="/keras_nlp/" role="tab" aria-selected="">KerasNLP: Natural Language Workflows</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/generative/'>Generative Deep Learning</a> / Neural Style Transfer with AdaIN </div> <div class='k-content'> <h1 id="neural-style-transfer-with-adain">Neural Style Transfer with AdaIN</h1> Author: <a href="https://twitter.com/arig23498">Aritra Roy Gosthipaty</a>, <a href="https://twitter.com/ritwik_raha">Ritwik Raha</a> Date created: 2021/11/08 Last modified: 2021/11/08 <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/adain.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/generative/adain.py">GitHub source</a> Description: Neural Style Transfer with Adaptive Instance Normalization. <h1 id="introduction">Introduction</h1> <a href="https://www.tensorflow.org/tutorials/generative/style_transfer">Neural Style Transfer</a> is the process of transferring the style of one image onto the content of another. This was first introduced in the seminal paper <a href="https://arxiv.org/abs/1508.06576">"A Neural Algorithm of Artistic Style"</a> by Gatys et al. A major limitation of the technique proposed in this work is in its runtime, as the algorithm uses a slow iterative optimization process. Follow-up papers that introduced <a href="https://arxiv.org/abs/1502.03167">Batch Normalization</a>, <a href="https://arxiv.org/abs/1701.02096">Instance Normalization</a> and <a href="https://arxiv.org/abs/1610.07629">Conditional Instance Normalization</a> allowed Style Transfer to be performed in new ways, no longer requiring a slow iterative process. Following these papers, the authors Xun Huang and Serge Belongie propose <a href="https://arxiv.org/abs/1703.06868">Adaptive Instance Normalization</a> (AdaIN), which allows arbitrary style transfer in real time. In this example we implement Adaptive Instance Normalization for Neural Style Transfer. We show in the below figure the output of our AdaIN model trained for only 30 epochs. <img alt="Style transfer sample gallery" src="https://i.imgur.com/zDjDuea.png" /> You can also try out the model with your own images with this <a href="https://huggingface.co/spaces/ariG23498/nst">Hugging Face demo</a>. <h1 id="setup">Setup</h1> We begin with importing the necessary packages. We also set the seed for reproducibility. The global variables are hyperparameters which we can change as we like. <div class="codehilite"><pre><code>import os import numpy as np import tensorflow as tf from tensorflow import keras import matplotlib.pyplot as plt import tensorflow_datasets as tfds from tensorflow.keras import layers # Defining the global variables. IMAGE_SIZE = (224, 224) BATCH_SIZE = 64 # Training for single epoch for time constraint. # Please use atleast 30 epochs to see good results. EPOCHS = 1 AUTOTUNE = tf.data.AUTOTUNE </code></pre></div> <hr /> <h2 id="style-transfer-sample-gallery">Style transfer sample gallery</h2> For Neural Style Transfer we need style images and content images. In this example we will use the <a href="https://www.kaggle.com/ikarus777/best-artworks-of-all-time">Best Artworks of All Time</a> as our style dataset and <a href="https://www.tensorflow.org/datasets/catalog/voc">Pascal VOC</a> as our content dataset. This is a deviation from the original paper implementation by the authors, where they use <a href="https://paperswithcode.com/dataset/wikiart">WIKI-Art</a> as style and <a href="https://cocodataset.org/#home">MSCOCO</a> as content datasets respectively. We do this to create a minimal yet reproducible example. <hr /> <h2 id="downloading-the-dataset-from-kaggle">Downloading the dataset from Kaggle</h2> The <a href="https://www.kaggle.com/ikarus777/best-artworks-of-all-time">Best Artworks of All Time</a> dataset is hosted on Kaggle and one can easily download it in Colab by following these steps: <ul> <li>Follow the instructions <a href="https://github.com/Kaggle/kaggle-api">here</a> in order to obtain your Kaggle API keys in case you don't have them.</li> <li>Use the following command to upload the Kaggle API keys.</li> </ul> <div class="codehilite"><pre><code>from google.colab import files files.upload() </code></pre></div> <ul> <li>Use the following commands to move the API keys to the proper directory and download the dataset.</li> </ul> <div class="codehilite"><pre><code>$ mkdir ~/.kaggle $ cp kaggle.json ~/.kaggle/ $ chmod 600 ~/.kaggle/kaggle.json $ kaggle datasets download ikarus777/best-artworks-of-all-time $ unzip -qq best-artworks-of-all-time.zip $ rm -rf images $ mv resized artwork $ rm best-artworks-of-all-time.zip artists.csv </code></pre></div> <hr /> <h2 id="tfdata"><a href="https://www.tensorflow.org/api_docs/python/tf/data"><code>tf.data</code></a> pipeline</h2> In this section, we will build the <a href="https://www.tensorflow.org/api_docs/python/tf/data"><code>tf.data</code></a> pipeline for the project. For the style dataset, we decode, convert and resize the images from the folder. For the content images we are already presented with a <a href="https://www.tensorflow.org/api_docs/python/tf/data"><code>tf.data</code></a> dataset as we use the <code>tfds</code> module. After we have our style and content data pipeline ready, we zip the two together to obtain the data pipeline that our model will consume. <div class="codehilite"><pre><code>def decode_and_resize(image_path): """Decodes and resizes an image from the image file path. Args: image_path: The image file path. Returns: A resized image. """ image = tf.io.read_file(image_path) image = tf.image.decode_jpeg(image, channels=3) image = tf.image.convert_image_dtype(image, dtype="float32") image = tf.image.resize(image, IMAGE_SIZE) return image def extract_image_from_voc(element): """Extracts image from the PascalVOC dataset. Args: element: A dictionary of data. Returns: A resized image. """ image = element["image"] image = tf.image.convert_image_dtype(image, dtype="float32") image = tf.image.resize(image, IMAGE_SIZE) return image # Get the image file paths for the style images. style_images = os.listdir("/content/artwork/resized") style_images = [os.path.join("/content/artwork/resized", path) for path in style_images] # split the style images in train, val and test total_style_images = len(style_images) train_style = style_images[: int(0.8 * total_style_images)] val_style = style_images[int(0.8 * total_style_images) : int(0.9 * total_style_images)] test_style = style_images[int(0.9 * total_style_images) :] # Build the style and content tf.data datasets. train_style_ds = ( tf.data.Dataset.from_tensor_slices(train_style) .map(decode_and_resize, num_parallel_calls=AUTOTUNE) .repeat() ) train_content_ds = tfds.load("voc", split="train").map(extract_image_from_voc).repeat() val_style_ds = ( tf.data.Dataset.from_tensor_slices(val_style) .map(decode_and_resize, num_parallel_calls=AUTOTUNE) .repeat() ) val_content_ds = ( tfds.load("voc", split="validation").map(extract_image_from_voc).repeat() ) test_style_ds = ( tf.data.Dataset.from_tensor_slices(test_style) .map(decode_and_resize, num_parallel_calls=AUTOTUNE) .repeat() ) test_content_ds = ( tfds.load("voc", split="test") .map(extract_image_from_voc, num_parallel_calls=AUTOTUNE) .repeat() ) # Zipping the style and content datasets. train_ds = ( tf.data.Dataset.zip((train_style_ds, train_content_ds)) .shuffle(BATCH_SIZE * 2) .batch(BATCH_SIZE) .prefetch(AUTOTUNE) ) val_ds = ( tf.data.Dataset.zip((val_style_ds, val_content_ds)) .shuffle(BATCH_SIZE * 2) .batch(BATCH_SIZE) .prefetch(AUTOTUNE) ) test_ds = ( tf.data.Dataset.zip((test_style_ds, test_content_ds)) .shuffle(BATCH_SIZE * 2) .batch(BATCH_SIZE) .prefetch(AUTOTUNE) ) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>[1mDownloading and preparing dataset voc/2007/4.0.0 (download: 868.85 MiB, generated: Unknown size, total: 868.85 MiB) to /root/tensorflow_datasets/voc/2007/4.0.0...[0m Dl Completed...: 0 url [00:00, ? url/s] Dl Size...: 0 MiB [00:00, ? MiB/s] Extraction completed...: 0 file [00:00, ? file/s] </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>0 examples [00:00, ? examples/s] Shuffling and writing examples to /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-test.tfrecord 0%| | 0/4952 [00:00<?, ? examples/s] 0 examples [00:00, ? examples/s] Shuffling and writing examples to /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-train.tfrecord 0%| | 0/2501 [00:00<?, ? examples/s] 0 examples [00:00, ? examples/s] Shuffling and writing examples to /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-validation.tfrecord 0%| | 0/2510 [00:00<?, ? examples/s] [1mDataset voc downloaded and prepared to /root/tensorflow_datasets/voc/2007/4.0.0. Subsequent calls will reuse this data.[0m </code></pre></div> </div> <hr /> <h2 id="visualizing-the-data">Visualizing the data</h2> It is always better to visualize the data before training. To ensure the correctness of our preprocessing pipeline, we visualize 10 samples from our dataset. <div class="codehilite"><pre><code>style, content = next(iter(train_ds)) fig, axes = plt.subplots(nrows=10, ncols=2, figsize=(5, 30)) [ax.axis("off") for ax in np.ravel(axes)] for (axis, style_image, content_image) in zip(axes, style[0:10], content[0:10]): (ax_style, ax_content) = axis ax_style.imshow(style_image) ax_style.set_title("Style Image") ax_content.imshow(content_image) ax_content.set_title("Content Image") </code></pre></div> <img alt="png" src="/img/examples/generative/adain/adain_8_0.png" /> <hr /> <h2 id="architecture">Architecture</h2> The style transfer network takes a content image and a style image as inputs and outputs the style transferred image. The authors of AdaIN propose a simple encoder-decoder structure for achieving this. <img alt="AdaIN architecture" src="https://i.imgur.com/JbIfoyE.png" /> The content image (<code>C</code>) and the style image (<code>S</code>) are both fed to the encoder networks. The output from these encoder networks (feature maps) are then fed to the AdaIN layer. The AdaIN layer computes a combined feature map. This feature map is then fed into a randomly initialized decoder network that serves as the generator for the neural style transferred image. <img alt="AdaIn equation" src="https://i.imgur.com/hqhcBQS.png" /> The style feature map (<code>fs</code>) and the content feature map (<code>fc</code>) are fed to the AdaIN layer. This layer produced the combined feature map <code>t</code>. The function <code>g</code> represents the decoder (generator) network. <h3 id="encoder">Encoder</h3> The encoder is a part of the pretrained (pretrained on <a href="https://www.image-net.org/">imagenet</a>) VGG19 model. We slice the model from the <code>block4-conv1</code> layer. The output layer is as suggested by the authors in their paper. <div class="codehilite"><pre><code>def get_encoder(): vgg19 = keras.applications.VGG19( include_top=False, weights="imagenet", input_shape=(*IMAGE_SIZE, 3), ) vgg19.trainable = False mini_vgg19 = keras.Model(vgg19.input, vgg19.get_layer("block4_conv1").output) inputs = layers.Input([*IMAGE_SIZE, 3]) mini_vgg19_out = mini_vgg19(inputs) return keras.Model(inputs, mini_vgg19_out, name="mini_vgg19") </code></pre></div> <h3 id="adaptive-instance-normalization">Adaptive Instance Normalization</h3> The AdaIN layer takes in the features of the content and style image. The layer can be defined via the following equation: <img alt="AdaIn formula" src="https://i.imgur.com/tWq3VKP.png" /> where <code>sigma</code> is the standard deviation and <code>mu</code> is the mean for the concerned variable. In the above equation the mean and variance of the content feature map <code>fc</code> is aligned with the mean and variance of the style feature maps <code>fs</code>. It is important to note that the AdaIN layer proposed by the authors uses no other parameters apart from mean and variance. The layer also does not have any trainable parameters. This is why we use a Python function instead of using a Keras layer. The function takes style and content feature maps, computes the mean and standard deviation of the images and returns the adaptive instance normalized feature map. <div class="codehilite"><pre><code>def get_mean_std(x, epsilon=1e-5): axes = [1, 2] # Compute the mean and standard deviation of a tensor. mean, variance = tf.nn.moments(x, axes=axes, keepdims=True) standard_deviation = tf.sqrt(variance + epsilon) return mean, standard_deviation def ada_in(style, content): """Computes the AdaIn feature map. Args: style: The style feature map. content: The content feature map. Returns: The AdaIN feature map. """ content_mean, content_std = get_mean_std(content) style_mean, style_std = get_mean_std(style) t = style_std * (content - content_mean) / content_std + style_mean return t </code></pre></div> <h3 id="decoder">Decoder</h3> The authors specify that the decoder network must mirror the encoder network. We have symmetrically inverted the encoder to build our decoder. We have used <code>UpSampling2D</code> layers to increase the spatial resolution of the feature maps. Note that the authors warn against using any normalization layer in the decoder network, and do indeed go on to show that including batch normalization or instance normalization hurts the performance of the overall network. This is the only portion of the entire architecture that is trainable. <div class="codehilite"><pre><code>def get_decoder(): config = {"kernel_size": 3, "strides": 1, "padding": "same", "activation": "relu"} decoder = keras.Sequential( [ layers.InputLayer((None, None, 512)), layers.Conv2D(filters=512, **config), layers.UpSampling2D(), layers.Conv2D(filters=256, **config), layers.Conv2D(filters=256, **config), layers.Conv2D(filters=256, **config), layers.Conv2D(filters=256, **config), layers.UpSampling2D(), layers.Conv2D(filters=128, **config), layers.Conv2D(filters=128, **config), layers.UpSampling2D(), layers.Conv2D(filters=64, **config), layers.Conv2D( filters=3, kernel_size=3, strides=1, padding="same", activation="sigmoid", ), ] ) return decoder </code></pre></div> <h3 id="loss-functions">Loss functions</h3> Here we build the loss functions for the neural style transfer model. The authors propose to use a pretrained VGG-19 to compute the loss function of the network. It is important to keep in mind that this will be used for training only the decoder network. The total loss (<code>Lt</code>) is a weighted combination of content loss (<code>Lc</code>) and style loss (<code>Ls</code>). The <code>lambda</code> term is used to vary the amount of style transferred. <img alt="The total loss" src="https://i.imgur.com/Q5y1jUM.png" /> <h3 id="content-loss">Content Loss</h3> This is the Euclidean distance between the content image features and the features of the neural style transferred image. <img alt="The content loss" src="https://i.imgur.com/dZ0uD0N.png" /> Here the authors propose to use the output from the AdaIn layer <code>t</code> as the content target rather than using features of the original image as target. This is done to speed up convergence. <h3 id="style-loss">Style Loss</h3> Rather than using the more commonly used <a href="https://mathworld.wolfram.com/GramMatrix.html">Gram Matrix</a>, the authors propose to compute the difference between the statistical features (mean and variance) which makes it conceptually cleaner. This can be easily visualized via the following equation: <img alt="The style loss" src="https://i.imgur.com/Ctclhn3.png" /> where <code>theta</code> denotes the layers in VGG-19 used to compute the loss. In this case this corresponds to: <ul> <li><code>block1_conv1</code></li> <li><code>block1_conv2</code></li> <li><code>block1_conv3</code></li> <li><code>block1_conv4</code></li> </ul> <div class="codehilite"><pre><code>def get_loss_net(): vgg19 = keras.applications.VGG19( include_top=False, weights="imagenet", input_shape=(*IMAGE_SIZE, 3) ) vgg19.trainable = False layer_names = ["block1_conv1", "block2_conv1", "block3_conv1", "block4_conv1"] outputs = [vgg19.get_layer(name).output for name in layer_names] mini_vgg19 = keras.Model(vgg19.input, outputs) inputs = layers.Input([*IMAGE_SIZE, 3]) mini_vgg19_out = mini_vgg19(inputs) return keras.Model(inputs, mini_vgg19_out, name="loss_net") </code></pre></div> <hr /> <h2 id="neural-style-transfer">Neural Style Transfer</h2> This is the trainer module. We wrap the encoder and decoder inside a <a href="/api/models/model#model-class"><code>tf.keras.Model</code></a> subclass. This allows us to customize what happens in the <code>model.fit()</code> loop. <div class="codehilite"><pre><code>class NeuralStyleTransfer(tf.keras.Model): def __init__(self, encoder, decoder, loss_net, style_weight, **kwargs): super().__init__(**kwargs) self.encoder = encoder self.decoder = decoder self.loss_net = loss_net self.style_weight = style_weight def compile(self, optimizer, loss_fn): super().compile() self.optimizer = optimizer self.loss_fn = loss_fn self.style_loss_tracker = keras.metrics.Mean(name="style_loss") self.content_loss_tracker = keras.metrics.Mean(name="content_loss") self.total_loss_tracker = keras.metrics.Mean(name="total_loss") def train_step(self, inputs): style, content = inputs # Initialize the content and style loss. loss_content = 0.0 loss_style = 0.0 with tf.GradientTape() as tape: # Encode the style and content image. style_encoded = self.encoder(style) content_encoded = self.encoder(content) # Compute the AdaIN target feature maps. t = ada_in(style=style_encoded, content=content_encoded) # Generate the neural style transferred image. reconstructed_image = self.decoder(t) # Compute the losses. reconstructed_vgg_features = self.loss_net(reconstructed_image) style_vgg_features = self.loss_net(style) loss_content = self.loss_fn(t, reconstructed_vgg_features[-1]) for inp, out in zip(style_vgg_features, reconstructed_vgg_features): mean_inp, std_inp = get_mean_std(inp) mean_out, std_out = get_mean_std(out) loss_style += self.loss_fn(mean_inp, mean_out) + self.loss_fn( std_inp, std_out ) loss_style = self.style_weight * loss_style total_loss = loss_content + loss_style # Compute gradients and optimize the decoder. trainable_vars = self.decoder.trainable_variables gradients = tape.gradient(total_loss, trainable_vars) self.optimizer.apply_gradients(zip(gradients, trainable_vars)) # Update the trackers. self.style_loss_tracker.update_state(loss_style) self.content_loss_tracker.update_state(loss_content) self.total_loss_tracker.update_state(total_loss) return { "style_loss": self.style_loss_tracker.result(), "content_loss": self.content_loss_tracker.result(), "total_loss": self.total_loss_tracker.result(), } def test_step(self, inputs): style, content = inputs # Initialize the content and style loss. loss_content = 0.0 loss_style = 0.0 # Encode the style and content image. style_encoded = self.encoder(style) content_encoded = self.encoder(content) # Compute the AdaIN target feature maps. t = ada_in(style=style_encoded, content=content_encoded) # Generate the neural style transferred image. reconstructed_image = self.decoder(t) # Compute the losses. recons_vgg_features = self.loss_net(reconstructed_image) style_vgg_features = self.loss_net(style) loss_content = self.loss_fn(t, recons_vgg_features[-1]) for inp, out in zip(style_vgg_features, recons_vgg_features): mean_inp, std_inp = get_mean_std(inp) mean_out, std_out = get_mean_std(out) loss_style += self.loss_fn(mean_inp, mean_out) + self.loss_fn( std_inp, std_out ) loss_style = self.style_weight * loss_style total_loss = loss_content + loss_style # Update the trackers. self.style_loss_tracker.update_state(loss_style) self.content_loss_tracker.update_state(loss_content) self.total_loss_tracker.update_state(total_loss) return { "style_loss": self.style_loss_tracker.result(), "content_loss": self.content_loss_tracker.result(), "total_loss": self.total_loss_tracker.result(), } @property def metrics(self): return [ self.style_loss_tracker, self.content_loss_tracker, self.total_loss_tracker, ] </code></pre></div> <hr /> <h2 id="train-monitor-callback">Train Monitor callback</h2> This callback is used to visualize the style transfer output of the model at the end of each epoch. The objective of style transfer cannot be quantified properly, and is to be subjectively evaluated by an audience. For this reason, visualization is a key aspect of evaluating the model. <div class="codehilite"><pre><code>test_style, test_content = next(iter(test_ds)) class TrainMonitor(tf.keras.callbacks.Callback): def on_epoch_end(self, epoch, logs=None): # Encode the style and content image. test_style_encoded = self.model.encoder(test_style) test_content_encoded = self.model.encoder(test_content) # Compute the AdaIN features. test_t = ada_in(style=test_style_encoded, content=test_content_encoded) test_reconstructed_image = self.model.decoder(test_t) # Plot the Style, Content and the NST image. fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(20, 5)) ax[0].imshow(tf.keras.utils.array_to_img(test_style[0])) ax[0].set_title(f"Style: {epoch:03d}") ax[1].imshow(tf.keras.utils.array_to_img(test_content[0])) ax[1].set_title(f"Content: {epoch:03d}") ax[2].imshow( tf.keras.utils.array_to_img(test_reconstructed_image[0]) ) ax[2].set_title(f"NST: {epoch:03d}") plt.show() plt.close() </code></pre></div> <hr /> <h2 id="train-the-model">Train the model</h2> In this section, we define the optimizer, the loss function, and the trainer module. We compile the trainer module with the optimizer and the loss function and then train it. Note: We train the model for a single epoch for time constraints, but we will need to train is for atleast 30 epochs to see good results. <div class="codehilite"><pre><code>optimizer = keras.optimizers.Adam(learning_rate=1e-5) loss_fn = keras.losses.MeanSquaredError() encoder = get_encoder() loss_net = get_loss_net() decoder = get_decoder() model = NeuralStyleTransfer( encoder=encoder, decoder=decoder, loss_net=loss_net, style_weight=4.0 ) model.compile(optimizer=optimizer, loss_fn=loss_fn) history = model.fit( train_ds, epochs=EPOCHS, steps_per_epoch=50, validation_data=val_ds, validation_steps=50, callbacks=[TrainMonitor()], ) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/vgg19/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5 80142336/80134624 [==============================] - 1s 0us/step 80150528/80134624 [==============================] - 1s 0us/step 50/50 [==============================] - ETA: 0s - style_loss: 213.1439 - content_loss: 141.1564 - total_loss: 354.3002 </code></pre></div> </div> <img alt="png" src="/img/examples/generative/adain/adain_23_1.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>50/50 [==============================] - 124s 2s/step - style_loss: 213.1439 - content_loss: 141.1564 - total_loss: 354.3002 - val_style_loss: 167.0819 - val_content_loss: 129.0497 - val_total_loss: 296.1316 </code></pre></div> </div> <hr /> <h2 id="inference">Inference</h2> After we train the model, we now need to run inference with it. We will pass arbitrary content and style images from the test dataset and take a look at the output images. NOTE: To try out the model on your own images, you can use this <a href="https://huggingface.co/spaces/ariG23498/nst">Hugging Face demo</a>. <div class="codehilite"><pre><code>for style, content in test_ds.take(1): style_encoded = model.encoder(style) content_encoded = model.encoder(content) t = ada_in(style=style_encoded, content=content_encoded) reconstructed_image = model.decoder(t) fig, axes = plt.subplots(nrows=10, ncols=3, figsize=(10, 30)) [ax.axis("off") for ax in np.ravel(axes)] for axis, style_image, content_image, reconstructed_image in zip( axes, style[0:10], content[0:10], reconstructed_image[0:10] ): (ax_style, ax_content, ax_reconstructed) = axis ax_style.imshow(style_image) ax_style.set_title("Style Image") ax_content.imshow(content_image) ax_content.set_title("Content Image") ax_reconstructed.imshow(reconstructed_image) ax_reconstructed.set_title("NST Image") </code></pre></div> <img alt="png" src="/img/examples/generative/adain/adain_25_0.png" /> <hr /> <h2 id="conclusion">Conclusion</h2> Adaptive Instance Normalization allows arbitrary style transfer in real time. It is also important to note that the novel proposition of the authors is to achieve this only by aligning the statistical features (mean and standard deviation) of the style and the content images. Note: AdaIN also serves as the base for <a href="https://arxiv.org/abs/1812.04948">Style-GANs</a>. <hr /> <h2 id="reference">Reference</h2> <ul> <li><a href="https://github.com/ftokarev/tf-adain">TF implementation</a></li> </ul> <hr /> <h2 id="acknowledgement">Acknowledgement</h2> We thank <a href="https://lukewood.xyz">Luke Wood</a> for his detailed review. </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#neural-style-transfer-with-adain'>Neural Style Transfer with AdaIN</a> </div> <div class='k-outline-depth-1'> <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-1'> <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#style-transfer-sample-gallery'>Style transfer sample gallery</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#downloading-the-dataset-from-kaggle'>Downloading the dataset from Kaggle</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#tfdata-pipeline'><code>tf.data</code> pipeline</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#visualizing-the-data'>Visualizing the data</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#architecture'>Architecture</a> </div> <div class='k-outline-depth-3'> <a href='#encoder'>Encoder</a> </div> <div class='k-outline-depth-3'> <a href='#adaptive-instance-normalization'>Adaptive Instance Normalization</a> </div> <div class='k-outline-depth-3'> <a href='#decoder'>Decoder</a> </div> <div class='k-outline-depth-3'> <a href='#loss-functions'>Loss functions</a> </div> <div class='k-outline-depth-3'> <a href='#content-loss'>Content Loss</a> </div> <div class='k-outline-depth-3'> <a href='#style-loss'>Style Loss</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#neural-style-transfer'>Neural Style Transfer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#train-monitor-callback'>Train Monitor callback</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#train-the-model'>Train the model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#inference'>Inference</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#conclusion'>Conclusion</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#reference'>Reference</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#acknowledgement'>Acknowledgement</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>