diff options
Diffstat (limited to 'losses.py')
| -rw-r--r-- | losses.py | 13 |
1 files changed, 6 insertions, 7 deletions
@@ -1,6 +1,6 @@ import tensorflow as tf from tensorflow.keras.layers import \ - Conv2D, MaxPool2D, Dropout, Flatten, Dense + Conv2D, MaxPool2D, Dropout, Flatten, Dense, AveragePool2D import numpy as np import hyperparameters as hp @@ -21,13 +21,13 @@ class YourModel(tf.keras.Model): activation="relu", name="block1_conv1"), Conv2D(64, 3, 1, padding="same", activation="relu", name="block1_conv2"), - MaxPool2D(2, name="block1_pool"), + AveragePool2D(2, name="block1_pool"), # Block 2 Conv2D(128, 3, 1, padding="same", activation="relu", name="block2_conv1"), Conv2D(128, 3, 1, padding="same", activation="relu", name="block2_conv2"), - MaxPool2D(2, name="block2_pool"), + AveragePool2D(2, name="block2_pool"), # Block 3 Conv2D(256, 3, 1, padding="same", activation="relu", name="block3_conv1"), @@ -35,7 +35,7 @@ class YourModel(tf.keras.Model): activation="relu", name="block3_conv2"), Conv2D(256, 3, 1, padding="same", activation="relu", name="block3_conv3"), - MaxPool2D(2, name="block3_pool"), + AveragePool2D(2, name="block3_pool"), # Block 4 Conv2D(512, 3, 1, padding="same", activation="relu", name="block4_conv1"), @@ -43,7 +43,7 @@ class YourModel(tf.keras.Model): activation="relu", name="block4_conv2"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block4_conv3"), - MaxPool2D(2, name="block4_pool"), + AveragePool2D(2, name="block4_pool"), # Block 5 Conv2D(512, 3, 1, padding="same", activation="relu", name="block5_conv1"), @@ -51,7 +51,7 @@ class YourModel(tf.keras.Model): activation="relu", name="block5_conv2"), Conv2D(512, 3, 1, padding="same", activation="relu", name="block5_conv3"), - MaxPool2D(2, name="block5_pool"), + AveragePool2D(2, name="block5_pool"), ] self.head = [ @@ -131,4 +131,3 @@ class YourModel(tf.keras.Model): L_style += self.layer_loss(art_layers, input_layers, layer) return L_style - |