diff options
Diffstat (limited to 'losses.py')
| -rw-r--r-- | losses.py | 19 |
1 files changed, 16 insertions, 3 deletions
@@ -12,14 +12,19 @@ class YourModel(tf.keras.Model): def __init__(self, content_image, style_image): #normalize these images to float values super(YourModel, self).__init__() - self.content_image = content_image + self.content_image = transform.resize(content_image, np.shape(style_image), anti_aliasing=True) + self.content_image = np.expand_dims(self.content_image, axis=0) #perhaps consider cropping to avoid distortion self.style_image = transform.resize(style_image, np.shape(style_image), anti_aliasing=True) + self.style_image = np.expand_dims(self.style_image, axis=0) self.x = tf.Variable(tf.random.uniform(np.shape(content_image)), trainable=True) + self.x = tf.expand_dims(self.x, axis=0) self.alpha = hp.alpha self.beta = hp.beta + print(self.content_image.shape, self.style_image.shape) + self.optimizer = tf.keras.optimizers.RMSprop(learning_rate=hp.learning_rate, momentum=hp.momentum) self.vgg16 = [ @@ -60,16 +65,19 @@ class YourModel(tf.keras.Model): activation="relu", name="block5_conv3"), AveragePooling2D(2, name="block5_pool"), ] - for layer in self.vgg16: - layer.trainable = False + # for layer in self.vgg16: + # layer.trainable = False self.indexed_layers = [layer for layer in self.vgg16 if layer.name == "conv1"] + print(self.indexed_layers) self.desired = [layer.name for layer in self.vgg16 if layer.name == "conv1"] def call(self, x): layers = [] for layer in self.vgg16: # pass the x through + print('this is', x.shape) + print(layer.name) x = layer(x) # print("Sotech117 is so so sus") @@ -87,10 +95,12 @@ class YourModel(tf.keras.Model): content_l = self.content_loss(photo_layers, input_layers) style_l = self.style_loss(art_layers, input_layers) # Equation 7 + print(content_l, style_l) 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)) + print(L_content) return L_content def layer_loss(art_layers, input_layers, layer): @@ -123,12 +133,15 @@ class YourModel(tf.keras.Model): L_style = 0 for layer in self.indexed_layers: L_style += self.layer_loss(art_layers, input_layers, layer) + print('this is style loss',L_style) return L_style def train_step(self): with tf.GradientTape() as tape: loss = self.loss_fn(self.content_image, self.style_image, self.x) + print(loss) gradients = tape.gradient(loss, self.x) + print(gradients) self.optimizer.apply_gradients(zip(gradients, self.x)) |