diff options
Diffstat (limited to 'losses.py')
| -rw-r--r-- | losses.py | 101 |
1 files changed, 59 insertions, 42 deletions
@@ -1,3 +1,4 @@ +from tkinter.tix import TCL_FILE_EVENTS import tensorflow as tf from tensorflow.keras.layers import \ Conv2D, MaxPool2D, Dropout, Flatten, Dense, AveragePooling2D @@ -18,8 +19,7 @@ class YourModel(tf.keras.Model): #perhaps consider cropping to avoid distortion self.style_image = transform.resize(style_image, tf.shape(style_image), anti_aliasing=True) self.style_image = tf.expand_dims(self.style_image, axis=0) - self.x = tf.Variable(tf.random.uniform(tf.shape(content_image)), trainable=True) - self.x = tf.expand_dims(self.x, axis=0) + self.x = tf.Variable(tf.expand_dims(tf.random.uniform(tf.shape(content_image)), axis=0), trainable=True) self.alpha = hp.alpha self.beta = hp.beta @@ -75,7 +75,7 @@ class YourModel(tf.keras.Model): self.layer_to_filters = {layer.name: layer.filters for layer in self.vgg16 if "conv" in layer.name} def call(self, x): - layers = np.empty(0) + layers = [] for layer in self.vgg16: # pass the x through x = layer(x) @@ -83,9 +83,9 @@ class YourModel(tf.keras.Model): # save the output of each layer if it is in the desired list if layer.name in self.desired: - layers = np.append(layers, x) + layers.append(x) - return x, tf.Variable(layers, dtype=np.float32) + return x, layers def loss_fn(self, p, a, x): _, photo_layers = self.call(p) @@ -98,54 +98,71 @@ class YourModel(tf.keras.Model): return (self.alpha * content_l) + (self.beta * style_l) def content_loss(self, photo_layers, input_layers): - L_content = tf.reduce_mean(tf.square(photo_layers - input_layers)) + L_content = tf.constant(0.0) + for i in range(len(photo_layers)): + pl = photo_layers[i] + il = input_layers[i] + L_content = tf.math.add(L_content, tf.reduce_mean(tf.square(pl - il))) + print('content loss', L_content) return L_content - def layer_loss(self, art_layers, input_layers, layer): - # vectorize the art_layers - art_layers = tf.reshape(art_layers, (-1, art_layers.shape[-1])) - # vectorize the input_layers - input_layers = tf.reshape(input_layers, (-1, input_layers.shape[-1])) - - # get the gram matrix - input_dim = input_layers.shape[0] - G = np.zeros((input_dim, input_dim)) - - for i in range(input_dim): - for j in range(input_dim): - k = np.dot(input_layers[i], art_layers[j]) - G[i,j] = k - G = tf.Variable(G, dtype=np.float32) - - # get the loss per each lateral layer - # N depends on # of filters in the layer, M depends on hight and width of feature map - M_l = art_layers.shape[0] * art_layers.shape[1] - - # layer.filters might not work - E_l = 1/4 * (self.layer_to_filters[layer.name]**(-2)) * (M_l**(-2)) * tf.reduce_sum(tf.square(G - input_layers)) + def layer_loss(self, art_layer, input_layer): + # vectorize the art_layers + art_layer = tf.reshape(art_layer, (-1, art_layer.shape[-1])) + # # vectorize the input_layers + input_layer = tf.reshape(input_layer, (-1, input_layer.shape[-1])) + # get the gram matrices + G_l = tf.matmul(tf.transpose(input_layer), input_layer) + A_l = tf.matmul(tf.transpose(art_layer), art_layer) + + + # vals = [] + # for i in range(input_dim): + # vals_i = [] + # for j in range(input_dim): + # il = tf.reshape(input_layers[i], [-1]) + # al = tf.reshape(art_layers[j], [-1]) + # k = tf.reduce_sum(tf.multiply(il, al)) + # vals_i.append(k) + # vals.append(tf.stack(vals_i)) + # G = tf.stack(vals) + + # get the loss per each lateral layer + # N depends on # of filters in the layer, M depends on hight and width of feature map + M_l = art_layer.shape[0] + N_l = art_layer.shape[1] + + # layer.filters might not work + E_l = 1/4 * (M_l**(-2)) *(N_l**(-2)) * tf.reduce_sum(tf.square(G_l - A_l)) - # while Sotech is botty: - # Jayson_tatum.tear_acl() - # return ("this is just another day") - print('Layer loss', E_l) - return E_l + # while Sotech is botty: + # Jayson_tatum.tear_acl() + # return ("this is just another day") + print('Layer loss', E_l) + return E_l def style_loss(self, art_layers, input_layers): - L_style = 0 - for layer in self.indexed_layers: - L_style += self.layer_loss(art_layers, input_layers, layer) + L_style = tf.constant(0.0) + for i in range(len(art_layers)): + art_layer = art_layers[i] + input_layer = input_layers[i] + L_style = tf.math.add(L_style, self.layer_loss(art_layer, input_layer)) print('style loss', L_style) return L_style def train_step(self): - with tf.GradientTape() as tape: + 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) - print('loss', loss) - print('self.x', self.x) - gradients = tape.gradient(loss, self.x) - print('gradients', gradients) - self.optimizer.apply_gradients(zip(gradients, self.x)) + #print('loss', loss) + #print('self.x', self.x) + gradients = tape.gradient(loss, [self.x]) + #print('gradients', gradients) + print(self.x.shape) + print(type(self.x)) + print(type(gradients)) + self.optimizer.apply_gradients(zip(gradients, [self.x])) |