PyTorch: Style Transfer
By: Karthik Bhaskar
Tags: Artificial Intelligence, Machine Learning, Deep Learning, Data Science, Art

In this article, we will recreate the style transfer method that is outlined in the paper, Image Style Transfer Using Convolutional Neural Networks, in PyTorch.

Image style transfer is a technique that aims to render the content of one image with the style of another, which is essential and exciting for both practical and scientific reasons. The style transfer techniques are widely used in image processing applications such as mobile camera filters and creative image generation.

In the article, we are going to use the pre-trained 19 layer VGG (Visual Geometry Group) Network to achieve style transfer task. VGG network consists of a series of convolutional, pooling, and fully connected layers. In the below image, we name the convolutional layer by stack and the order of the stack. For example, conv_1_1 represents the first convolutional layer in the first stack; conv_2_1 represents the first convolutional layer in the second stack. In the architecture shown below, the deepest convolutional layer in the network is conv_5_4.

For the Style transfer task, first, we need to have two images.

Content Image — represents the image that you want to style Style Image — represents the style, colors, and textures that you want to use on the content image.

We use these content image and style images to create a new target image, which has both the properties of style and content image. Here, we have a beautiful picture of Gal Gadot as a content image and an abstract art design as a style image. We want to use the style image on Gal Gadot’s picture to obtain an artistic sketched portrait of Gal Gadot.

Let’s Start the Magic in PyTorch

We will use a pre-trained VGG19 Net to extract content or style features. Then we will formalize the notion of content and style losses and apply those to iteratively update our target image until we get a result that we want. Let’s import the necessary resources for our model.

# import resources

%matplotlib inline from PIL import Image

from io import BytesIO

import matplotlib.pyplot as plt

import numpy as np import torch

import torch.optim as optim

import requests

from torchvision import transforms, models

Load Pre-trained VGG Model

PyTorch’s VGG19 pre-trained model has two parts. vgg19.features contain convolutional and pooling layers, and vgg19.classifier has 3 fully connected classifiers. We only need vgg19.features to extract the content and style features of the image, so we will load these and freeze the weights.

# get the "features" portion of VGG19 (we will not need the "classifier" portion) vgg = models.vgg19(pretrained=True).features # freeze all VGG parameters since we're only optimizing the target image for param in vgg.parameters():

param.requires_grad_(False) # move the model to GPU, if available device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

vgg.to(device)

Below are the helper functions to load the content and style images and to convert the images to normalized tensors. These functions will resize the content and style images of the same size. Give the appropriate image location to these functions.

Features and Gram Matrix

To get the content and style representations of an image, we have to pass an image forward through the VGG19 network until we get to the desired layer(s) and then get the output from that layer. In the Style Transfer paper, they used conv1_1 (layer 0), conv2_1 (layer 5), conv_3_1 (layer 10), conv4_1 (layer 19) and conv_5_1 (layer 28) for Style Representation and they used conv5_1 (layer 28) for Content Representation. The function get_features() below does this. We need something called Gram Matrix to get the style representations in the required format to train our model. The function gram_matrix() below does this. We use the get_features() function to get the content and style features, and then we calculate the gram matrices for each layer of our style representation. We now create a “target image” to combine the style and content representation together. We make a copy of our content image as a starting point and then iteratively change its style.

Individual Layer Style Weights

Below, we have the option to weight the style representation at each relevant layer. It’s suggested that you use a range between 0–1 to weight these layers. By weighting earlier layers (conv1_1 and conv2_1) more, you can expect to get larger style artifacts in your resulting target image.

Content and Style Weight

Just like in the paper, we define an alpha (content_weight) and a beta (style_weight). This ratio will affect how stylized your final image is. It’s recommended that you leave the content_weight = 1 and set the style_weight to achieve the ratio you want.

Training the Model

We will decide on a number of steps for which to update our target image, we are only changing our target image and nothing else about VGG19. I used 3000 steps in our training loop. Inside the iteration loop, we will calculate the content and style losses and update our target image. Content loss is the MSE between target and content features. Style loss is also calculated in a similar way, where we iterate through the layer mentioned in style_weights. Finally, we will create the total loss by adding style and content losses and weighting them with our specified alpha and beta values. The following snippet gives the Content Weight, Style Weight and the Training Loop.

The below are the results of our training printed out for every 400 steps.

Results of Every 400 Steps during Training

Finally, after 3000 steps, we have obtained an artistic sketched portrait of Gal Gadot.

Original Image vs Style Transfered Image

Please look at the code snippets for detailed understanding.

Reference: