TensorFlow Migration from 1.2 to 2.0: A Picture-Style Transfer Example

Artificial Intelligence
1

TensorFlow Migration from 1.2 to 2.0: A Picture-Style Transfer Example

In this article, we will explore the migration of a picture-style transfer example from TensorFlow 1.2 to 2.0. The example uses a convolutional neural network (CNN) to transfer the style of a famous artwork onto a photograph. We will discuss the changes made to the code and the improvements achieved by migrating to TensorFlow 2.0.

From the Boiler Worker to AI Expert

The original code used TensorFlow 1.x, which had a complex algorithm and structure. With the help of TensorFlow 2.0 and its new optimization features, we were able to simplify the code and achieve the same results with significantly reduced computational overhead.

Style Transfer

Style transfer refers to the technique of using neural networks to generate images that resemble the style of a particular artwork, while maintaining the content of the original photograph. This is achieved by training a CNN to extract the style features of the artwork and apply them to the photograph.

Fundamental Style Transfer

The fundamental style transfer algorithm is based on the principle of transfer paper, which uses a neural network to extract the style features of an image and apply them to another image. The algorithm uses a combination of convolutional and pooling layers to extract the style features, and then uses a fully connected layer to apply these features to the target image.

Source Code

The source code for the picture-style transfer example is shown below:

# Import necessary libraries
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import time
import functools
import time
from PIL import Image

# Set the drawing window parameters for image display
mpl.rcParams['figure.figsize'] = (13, 10)
mpl.rcParams['axes.grid'] = False

# Get the path of the local picture download, content_path is a real photo, style_path is art style picture
content_path = "1-content.jpg"
style_path = "1-style.jpg"

# Read a picture, and preprocessing
def load_img(path_to_img):
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_jpeg(img)
    img = tf.image.convert_image_dtype(img, tf.float32)
    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    long = max(shape)
    scale = max_dim / long
    new_shape = tf.cast(shape * scale, tf.int32)
    img = img[tf.newaxis, :]
    return img

# Read two pictures
content_image = load_img(content_path)
style_image = load_img(style_path)

# Define the content and style layers
content_layers = ['block5_conv2']
style_layers = ['block1_conv1', 'Block2_conv1', 'Block3_conv1', 'Block4_conv1', 'Block5_conv1']

# Define a tool function to extract the intermediate layer output
def vgg_layers(layer_names):
    """Creates a vgg model that returns a list of intermediate output values."""
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
    vgg.trainable = False
    outputs = [vgg.get_layer(name).output for name in layer_names]
    model = tf.keras.Model([vgg.input], outputs)
    return model

# Define the style matrix function
def gram_matrix(input_tensor):
    result = tf.linalg.einsum('bijc, bijd->bcd', input_tensor, input_tensor)
    input_shape = tf.shape(input_tensor)
    num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
    return result / num_locations

# Define the StyleContentModel class
class StyleContentModel(tf.keras.models.Model):
    def __init__(self, style_layers, content_layers):
        super(StyleContentModel, self).__init__()
        self.vgg = vgg_layers(style_layers + content_layers)
        self.style_layers = style_layers
        self.content_layers = content_layers
        self.num_style_layers = len(style_layers)

    def call(self, input):
        # Preprocess the input image
        vgg_input = input * 255.0
        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(vgg_input)

        # Extract the style and content outputs
        outputs = self.vgg(preprocessed_input)
        style_outputs, content_outputs = (outputs[:self.num_style_layers], outputs[self.num_style_layers:])

        # Calculate the style matrix
        style_outputs = [gram_matrix(style_output) for style_output in style_outputs]

        # Convert the content and style outputs to dictionaries
        content_dict = {content_name: value for content_name, value in zip(self.content_layers, content_outputs)}
        style_dict = {style_name: value for style_name, value in zip(self.style_layers, style_outputs)}

        # Return the content and style results
        return {'content': content_dict, 'style': style_dict}

# Create an instance of the StyleContentModel class
extractor = StyleContentModel(style_layers, content_layers)

# Set the style and content targets
style_targets = extractor(style_image)['style']
content_targets = extractor(content_image)['content']

# Define the clip_0_1 function
def clip_0_1(image):
    return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

# Define the style_content_loss function
def style_content_loss(outputs):
    style_outputs = outputs['style']
    content_outputs = outputs['content']
    style_loss = tf.add_n([tf.reduce_mean((style_outputs[name] - style_targets[name]) ** 2) for name in style_outputs.keys()])
    style_loss *= style_weight / len(style_layers)
    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name] - content_targets[name]) ** 2) for name in content_outputs.keys()])
    content_loss *= content_weight / len(content_layers)
    loss = style_loss + content_loss
    return loss

# Define the train_step function
@tf.function()
def train_step(image):
    with tf.GradientTape() as tape:
        # Extract the style and content outputs
        outputs = extractor(image)

        # Calculate the loss
        loss = style_content_loss(outputs)

        # Gradient descent
        grad = tape.gradient(loss, image)

        # Apply the gradients
        opt.apply_gradients([(grad, image)])

        # Update the image
        image.assign(clip_0_1(image))

# Train the model
start = time.time()
epochs = 10
steps_per_epoch = 100
step = 0
for n in range(epochs):
    for m in range(steps_per_epoch):
        step += 1
        train_step(image)
    print('.', end='')
print('')

# Save the resulting image
file_name = 'newart1.png'
mpl.image.imsave(file_name, image[0])

Changes Made to the Code

The code was modified to use TensorFlow 2.0 and its new optimization features. The following changes were made:

  • The tf.keras.applications.VGG19 model was replaced with the tf.keras.applications.VGG19(include_top=False, weights='imagenet') model.
  • The tf.keras.layers.Conv2D and tf.keras.layers.MaxPooling2D layers were replaced with the tf.keras.layers.Conv2D and tf.keras.layers.AveragePooling2D layers.
  • The tf.keras.models.Model class was used to define the StyleContentModel class.
  • The tf.keras.optimizers.Adam optimizer was used to train the model.
  • The tf.clip_by_value function was used to clip the output of the model.

Improvements Achieved

The migration to TensorFlow 2.0 resulted in significant improvements in the code and the results. The following improvements were achieved:

  • The code was simplified and easier to understand.
  • The computational overhead was reduced.
  • The model was trained faster.
  • The results were improved.

Conclusion

The migration of the picture-style transfer example from TensorFlow 1.2 to 2.0 resulted in significant improvements in the code and the results. The use of TensorFlow 2.0 and its new optimization features simplified the code and reduced the computational overhead. The model was trained faster and the results were improved.