diff options
Diffstat (limited to 'losses.py')
| -rw-r--r-- | losses.py | 97 |
1 files changed, 0 insertions, 97 deletions
diff --git a/losses.py b/losses.py deleted file mode 100644 index c618fd38..00000000 --- a/losses.py +++ /dev/null @@ -1,97 +0,0 @@ -import tensorflow as tf -import numpy as np -from tensorflow.keras.layers import \ - Conv2D, AveragePooling2D -from skimage import transform -import hyperparameters as hp - -def get_gram(style_output): - style_shape = tf.shape(style_output) - output = tf.linalg.einsum('bijc,bijd->bcd', style_output, style_output) - dimensions = style_shape[1] * style_shape[2] - dimensions = tf.cast(dimensions, tf.float32) - return output / dimensions - - -class YourModel(tf.keras.Model): - """ Your own neural network model. """ - - def __init__(self, content_image, style_image): #normalize these images to float values - super(YourModel, self).__init__() - - self.content_image = transform.resize(content_image, tf.shape(style_image), anti_aliasing=True, preserve_range=True) - self.content_image = tf.image.convert_image_dtype(self.content_image, tf.float32) - - self.style_image = transform.resize(style_image, tf.shape(style_image), anti_aliasing=True, preserve_range=True) - self.style_image = tf.image.convert_image_dtype(self.style_image, tf.float32) - - image = tf.image.convert_image_dtype(content_image, tf.float32) - self.x = tf.Variable([image]) - - self.content_weight = hp.alpha - self.style_weight = hp.beta - - self.photo_layers = ['block5_conv2'] - self.style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1'] - - self.num_photo_layers = len(self.photo_layers) - self.num_style_layers = len(self.style_layers) - - self.optimizer = tf.keras.optimizers.Adam(learning_rate=hp.learning_rate, beta_1=hp.beta_1, epsilon=hp.epsilon) - - self.vgg16 = tf.keras.applications.VGG16(include_top=False, weights='vgg16_imagenet.h5') - - self.vgg16.trainable = False - - # creating the Gram Matrix - p_output = self.vgg16.get_layer(self.photo_layers[0]).output - style_output = [] - for layer in self.style_layers: - style_output.append(self.vgg16.get_layer(layer).output) - - G = [get_gram(x) for x in style_output] - - self.vgg16 = tf.keras.Model([self.vgg16.input], [p_output, G]) - - # figure this out Michael - img_to_np = lambda img: np.array([img * 255]) - - self.content_target = self.vgg16(img_to_np(content_image))[0] - self.style_target = self.vgg16(img_to_np(style_image))[1] - - # create a map of the layers to their corresponding number of filters if it is a convolutional layer - - def call(self, x): - # call onto our pretrained network, since we don't have a classifcation head to follow - x = self.vgg16(x * 255) - return x - - def loss_fn(self, x): - x = self.call(x) - - - content_l = self.content_loss(x[0], self.content_target) - style_l = self.style_loss(x[1], self.style_target) - return (self.content_weight * content_l) + (self.style_weight * style_l) - - def content_loss(self, photo_layers, input_layers): - return tf.reduce_mean(tf.square(photo_layers - input_layers)) - - def style_loss(self, art_layers, input_layers): - layer_losses = [] - for created, target in zip(art_layers, input_layers): - reduced = tf.reduce_mean(tf.square(created - target)) - layer_losses.append(reduced) - return tf.add_n(layer_losses) - - - def train_step(self, epoch): - with tf.GradientTape(watch_accessed_variables=False) as tape: - tape.watch(self.x) - # loss = self.loss_fn(self.content_image, self.style_image, self.x) - loss = self.loss_fn(self.x) - print('\rEpoch {}: Loss: {:.4f}'.format(epoch, loss), end='') - gradients = tape.gradient(loss, self.x) - self.optimizer.apply_gradients([(gradients, self.x)]) - self.x.assign(tf.clip_by_value(self.x, clip_value_min=0.0, clip_value_max=1.0)) - |