[TOC]
style_transfer.py
""" Implementation in TensorFlow of the paper
A Neural Algorithm of Artistic Style (Gatys et al., 2016)
Created by Chip Huyen ([email protected])
CS20: "TensorFlow for Deep Learning Research"
cs20.stanford.edu
For more details, please read the assignment handout:
https://docs.google.com/document/d/1FpueD-3mScnD0SJQDtwmOb1FrSwo1NGowkXzMwPoLH4/edit?usp=sharing
"""
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import time
import numpy as np
import tensorflow as tf
import load_vgg
import utils
def setup():
utils.safe_mkdir('checkpoints')
utils.safe_mkdir('outputs')
class StyleTransfer(object):
def __init__(self, content_img, style_img, img_width, img_height):
'''
img_width and img_height are the dimensions we expect from the generated image.
We will resize input content image and input style image to match this dimension.
Feel free to alter any hyperparameter here and see how it affects your training.
'''
self.img_width = img_width
self.img_height = img_height
self.content_img = utils.get_resized_image(content_img, img_width, img_height)
self.style_img = utils.get_resized_image(style_img, img_width, img_height)
self.initial_img = utils.generate_noise_image(self.content_img, img_width, img_height)
###############################
## TO DO
## create global step (gstep) and hyperparameters for the model
self.gstep = tf.Variable(0, dtype=tf.int32, trainable=False, name="global_step")
self.content_layer = 'conv4_2'
self.style_layers = ['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1', 'conv5_1']
# content_w, style_w: corresponding weights for content loss and style loss
self.content_w = 0.01
self.style_w = 1
# style_layer_w: weights for different style layers. deep layers have more weights
self.style_layer_w = [0.5, 1.0, 1.5, 3.0, 4.0]
# self.gstep = None # global step
self.lr = 0.001
###############################
def create_input(self):
'''
We will use one input_img as a placeholder for the content image,
style image, and generated image, because:
1. they have the same dimension
2. we have to extract the same set of features from them
We use a variable instead of a placeholder because we're, at the same time,
training the generated image to get the desirable result.
Note: image height corresponds to number of rows, not columns.
'''
with tf.variable_scope('input') as scope:
self.input_img = tf.get_variable('in_img',
shape=([1, self.img_height, self.img_width, 3]),
dtype=tf.float32,
initializer=tf.zeros_initializer())
def load_vgg(self):
'''
Load the saved model parameters of VGG-19, using the input_img
as the input to compute the output at each layer of vgg.
During training, VGG-19 mean-centered all images and found the mean pixels
to be [123.68, 116.779, 103.939] along RGB dimensions. We have to subtract
this mean from our images.
'''
self.vgg = load_vgg.VGG(self.input_img)
self.vgg.load()
self.content_img -= self.vgg.mean_pixels
self.style_img -= self.vgg.mean_pixels
def _content_loss(self, P, F):
''' Calculate the loss between the feature representation of the
content image and the generated image.
Inputs:
P: content representation of the content image
F: content representation of the generated image
Read the assignment handout for more details
Note: Don't use the coefficient 0.5 as defined in the paper.
Use the coefficient defined in the assignment handout.
'''
###############################
## TO DO
coef =1/(4*p[1]*p[2]*p[3])
self.content_loss = coef*tf.reduce_sum(tf.square(P-F))
###############################
def _gram_matrix(self, F, N, M):
""" Create and return the gram matrix for tensor F
Hint: you'll first have to reshape F
"""
###############################
## TO DO
F = tf.reshpae(F, shape=[M, N])
gram = tf.matmul(tf.transpose(F),F)
return gram
###############################
def _single_style_loss(self, a, g):
""" Calculate the style loss at a certain layer
Inputs:
a is the feature representation of the style image at that layer
g is the feature representation of the generated image at that layer
Output:
the style loss at a certain layer (which is E_l in the paper)
Hint: 1. you'll have to use the function _gram_matrix()
2. we'll use the same coefficient for style loss as in the paper
3. a and g are feature representation, not gram matrices
"""
###############################
## TO DO
N = a.shape[3]
M = a.shape[1]*a.shape[2]
a = tf.reshape(a, [M,N])
G = self._gram_matrix(g, N, M)
E = 1/(4*N**2*M**2)*tf.reduce_sum(tf.square(a-G))
return E
###############################
def _style_loss(self, A):
""" Calculate the total style loss as a weighted sum
of style losses at all style layers
Hint: you'll have to use _single_style_loss()
"""
###############################
## TO DO
# copy the answer
n_layers = len(A)
E = [self._single_style_loss(A[i], getattr(self.vgg, self.style_layers[i])) for i in range(n_layers)]
self.style_loss = sum([self.style_layer_w[i]*E[i] for i in range(n_layers)])
# for i, a in enumerate(A):
# self.style_loss +=
###############################
def losses(self):
with tf.variable_scope('losses') as scope:
with tf.Session() as sess:
# assign content image to the input variable
sess.run(self.input_img.assign(self.content_img))
gen_img_content = getattr(self.vgg, self.content_layer)
content_img_content = sess.run(gen_img_content)
self._content_loss(content_img_content, gen_img_content)
with tf.Session() as sess:
sess.run(self.input_img.assign(self.style_img))
style_layers = sess.run([getattr(self.vgg, layer) for layer in self.style_layers])
self._style_loss(style_layers)
##########################################
## TO DO: create total loss.
## Hint: don't forget the weights for the content loss and style loss
self.total_loss = self.content_w*self.content_loss+self.style_w*self.style_loss
##########################################
def optimize(self):
###############################
## TO DO: create optimizer
self.opt = tf.train.AdamOptimizer(self.lr).minimize(self.loss, global_step=self.gstep)
###############################
def create_summary(self):
###############################
## TO DO: create summaries for all the losses
## Hint: don't forget to merge them
tf.summary.scalar("style_loss", self.style_loss)
tf.summary.scalar("content_loss", self.content_loss)
tf.summary.scalar("total_loss", self.total_loss)
tf.summary.histogram("histogram loss", self.total_loss)
self.summary_op = tf.summary.merge_all()
###############################
def build(self):
self.create_input()
self.load_vgg()
self.losses()
self.optimize()
self.create_summary()
def train(self, n_iters):
skip_step = 1
with tf.Session() as sess:
###############################
## TO DO:
## 1. initialize your variables
## 2. create writer to write your grapp
###############################
sess.run(tf.gloabl_variables_initializer())
sess.run(self.input_img.assign(self.initial_img))
writer = tf.summary.FileWriter('./graphs/style_transfer', tf.get_default_graph())
###############################
## TO DO:
## 1. create a saver object
## 2. check if a checkpoint exists, restore the variables
##############################
utils.safe_mkdir('checkpoints')
utils.safe_mkdir('checkpoints/style_transfer')
# writer = tf.summary.FileWriter("")
saver = tf.train.saver()
ckpt = tf.train.get_checkpoint_state(os.path.dirname('checkpoints/convnet_layers/checkpoint'))
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
initial_step = self.gstep.eval()
start_time = time.time()
for index in range(initial_step, n_iters):
if index >= 5 and index < 20:
skip_step = 10
elif index >= 20:
skip_step = 20
sess.run(self.opt)
if (index + 1) % skip_step == 0:
###############################
## TO DO: obtain generated image, loss, and summary
gen_image, total_loss, summary = sess.run([self.input_img, self.total_loss, self.summary_op])
###############################
# try:
# while True:
# _, total_loss, summary = sess.run([self.opt, self.losses, self.summary_op])
# writer.add_summary(summary, global_step=index)
# add back the mean pixels we subtracted before
gen_image = gen_image + self.vgg.mean_pixels
writer.add_summary(summary, global_step=index)
print('Step {}\n Sum: {:5.1f}'.format(index + 1, np.sum(gen_image)))
print(' Loss: {:5.1f}'.format(total_loss))
print(' Took: {} seconds'.format(time.time() - start_time))
start_time = time.time()
filename = 'outputs/%d.png' % (index)
utils.save_image(filename, gen_image)
if (index + 1) % 20 == 0:
###############################
## TO DO: save the variables into a checkpoint
###############################
saver.save(sess, "checkpoints/style_transfer/style_transfer", index);
if __name__ == '__main__':
setup()
machine = StyleTransfer('content/van.jpg', 'styles/deadpool.jpg', 333, 250)
machine.build()
machine.train(300)
load_vgg.py
""" Load VGGNet weights needed for the implementation in TensorFlow
of the paper A Neural Algorithm of Artistic Style (Gatys et al., 2016)
Created by Chip Huyen ([email protected])
CS20: "TensorFlow for Deep Learning Research"
cs20.stanford.edu
For more details, please read the assignment handout:
https://docs.google.com/document/d/1FpueD-3mScnD0SJQDtwmOb1FrSwo1NGowkXzMwPoLH4/edit?usp=sharing
"""
import numpy as np
import scipy.io
import tensorflow as tf
import utils
# VGG-19 parameters file
VGG_DOWNLOAD_LINK = 'http://www.vlfeat.org/matconvnet/models/imagenet-vgg-verydeep-19.mat'
VGG_FILENAME = 'imagenet-vgg-verydeep-19.mat'
EXPECTED_BYTES = 534904783
class VGG(object):
def __init__(self, input_img):
utils.download(VGG_DOWNLOAD_LINK, VGG_FILENAME, EXPECTED_BYTES)
self.vgg_layers = scipy.io.loadmat(VGG_FILENAME)['layers']
self.input_img = input_img
self.mean_pixels = np.array([123.68, 116.779, 103.939]).reshape((1,1,1,3))
def _weights(self, layer_idx, expected_layer_name):
""" Return the weights and biases at layer_idx already trained by VGG
"""
W = self.vgg_layers[0][layer_idx][0][0][2][0][0]
b = self.vgg_layers[0][layer_idx][0][0][2][0][1]
layer_name = self.vgg_layers[0][layer_idx][0][0][0][0]
assert layer_name == expected_layer_name
return W, b.reshape(b.size)
def conv2d_relu(self, prev_layer, layer_idx, layer_name):
""" Create a convolution layer with RELU using the weights and
biases extracted from the VGG model at 'layer_idx'. You should use
the function _weights() defined above to extract weights and biases.
_weights() returns numpy arrays, so you have to convert them to TF tensors.
Don't forget to apply relu to the output from dthe convolution.
Inputs:
prev_layer: the output tensor from the previous layer
layer_idx: the index to current layer in vgg_layers
layer_name: the string that is the name of the current layer.
It's used to specify variable_scope.
Hint for choosing strides size:
for small images, you probably don't want to skip any pixel
"""
###############################
## TO DO
with tf.variable_scope(layer_name, reuse=tf.AUTO_REUSE) as scope:
in_channels = prev_layer.shape[-1]
kernel, bias = self._weights(layer_idx, layer_name)
# kernel = tf.convert_to_tensor(kernel)
# bias = tf.convert_to_tensor(bias)
kernel = tf.constant(W, name="weights")
bias = tf.constant(b, name="bias")
conv = tf.nn.conv2d(prev_layer, kernel, strides=[1,1,1,1], padding="SAME")
out = tf.nn.relu(conv + biases, name=scope.name)
###############################
setattr(self, layer_name, out)
def avgpool(self, prev_layer, layer_name):
""" Create the average pooling layer. The paper suggests that
average pooling works better than max pooling.
Input:
prev_layer: the output tensor from the previous layer
layer_name: the string that you want to name the layer.
It's used to specify variable_scope.
Hint for choosing strides and kszie: choose what you feel appropriate
"""
###############################
## TO DO
with tf.variable_scope(layer_name, reuse=tf.AUTO_REUSE) as scope:
out = tf.nn.average_pool(prev_layer, ksize=[1,2,2,1], strides=[1,2,2,1], padding="SAME")
###############################
setattr(self, layer_name, out)
def load(self):
self.conv2d_relu(self.input_img, 0, 'conv1_1')
self.conv2d_relu(self.conv1_1, 2, 'conv1_2')
self.avgpool(self.conv1_2, 'avgpool1')
self.conv2d_relu(self.avgpool1, 5, 'conv2_1')
self.conv2d_relu(self.conv2_1, 7, 'conv2_2')
self.avgpool(self.conv2_2, 'avgpool2')
self.conv2d_relu(self.avgpool2, 10, 'conv3_1')
self.conv2d_relu(self.conv3_1, 12, 'conv3_2')
self.conv2d_relu(self.conv3_2, 14, 'conv3_3')
self.conv2d_relu(self.conv3_3, 16, 'conv3_4')
self.avgpool(self.conv3_4, 'avgpool3')
self.conv2d_relu(self.avgpool3, 19, 'conv4_1')
self.conv2d_relu(self.conv4_1, 21, 'conv4_2')
self.conv2d_relu(self.conv4_2, 23, 'conv4_3')
self.conv2d_relu(self.conv4_3, 25, 'conv4_4')
self.avgpool(self.conv4_4, 'avgpool4')
self.conv2d_relu(self.avgpool4, 28, 'conv5_1')
self.conv2d_relu(self.conv5_1, 30, 'conv5_2')
self.conv2d_relu(self.conv5_2, 32, 'conv5_3')
self.conv2d_relu(self.conv5_3, 34, 'conv5_4')
self.avgpool(self.conv5_4, 'avgpool5')