Pulling all toguether to training

Author: eddydecena

.gitignore

@@ -250,4 +250,8 @@ $RECYCLE.BIN/
 # Windows shortcuts
 *.lnk
 
+
+examples
+model_checkpoints
+
 # End of https://www.toptal.com/developers/gitignore/api/visualstudiocode,macos,windows,linux,git,virtualenv,python

cycle/__init__.py

@@ -0,0 +1 @@
+from .model import CycleGAN

cycle/callbacks.py

@@ -1,23 +1,24 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from tensorflow.data import Dataset
-from tensorflow.keras.callbacks import CallBack
-from tensorflow.keras.preprocessing import array_to_img
+from tensorflow.keras.callbacks import Callback
+from tensorflow.keras.preprocessing.image import array_to_img
 
-class GANMonitor(CallBack):
-    def __init__(self, img_num: int):
+class GANMonitor(Callback):
+    def __init__(self, data: Dataset, img_num: int = 4):
         super(GANMonitor, self).__init__()
         self.img_num: int = img_num
+        self.data = data
 
-    def on_epoch_end(self, epoch: int, test_horses: Dataset, logs=None):
+    def on_epoch_end(self, epoch: int, logs=None):
         _, ax = plt.subplots(self.img_num, 2, figsize=(12, 12))
-        for i, img in enumerate(test_horses.take(self.img_num)):
-            output = self.model.gen_G(img)[0].numpy()
+        for i, img in enumerate(self.data.take(self.img_num)):
+            output = self.model.generator_G(img)[0]
             output = (output * 127.5 + 127.5).numpy().astype(np.uint8)
             img = (img[0] * 127.5 + 127.5).numpy().astype(np.uint8)
 
             output = array_to_img(output)
             img = array_to_img(img)
 
-            output.save(f'generated_img_{i}_{epoch}.png')
-            img.save('original_img_{i}_{epoch}.png')
+            output.save(f'./examples/generated_img_{i}_{epoch}.png')
+            img.save('./examples/original_img_{i}_{epoch}.png')

cycle/model_helper.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 from tensorflow.keras import layers
 from tensorflow_addons.layers import InstanceNormalization
 from tensorflow.keras.initializers import Initializer
@@ -8,14 +10,14 @@
 def residual_block(
     x: layers.Layer,
     activation: layers.Activation,
-    kernel_size: tuple[int] = (3, 3),
-    strides: tuple[int] = (1, 1),
+    kernel_size: Tuple[int] = (3, 3),
+    strides: Tuple[int] = (1, 1),
     padding: str = 'valid',
     kernel_initializer: Initializer = None,
     gamma_initializer: Initializer = None,
     use_bias: bool = False) -> layers.Layer:
 
-    dim: int =- x.shape[-1]
+    dim: int = x.shape[-1]
     input_tensor: layers.Layer = x
 
     x = ReflectionPadding2D()(input_tensor)
@@ -69,8 +71,8 @@ def upsample(
     x: layers.Layer,
     filters: int,
     activation: layers.Activation,
-    kernel_size: tuple[int] = (3, 3),
-    strides: tuple[int] = (2, 2),
+    kernel_size: Tuple[int] = (3, 3),
+    strides: Tuple[int] = (2, 2),
     padding: str = 'same',
     kernel_initializer: Initializer = None,
     gamma_initializer: Initializer = None,

train.py

@@ -3,14 +3,30 @@
 import numpy as np
 import matplotlib.pyplot as plt
 import tensorflow_datasets as tfds
+from tensorflow.data import Dataset
+from tensorflow.keras.optimizers import Adam
 from tensorflow.data.experimental import AUTOTUNE
+from tensorflow.keras.callbacks import ModelCheckpoint
+from tensorflow.config.experimental import set_memory_growth
+from tensorflow.config import list_physical_devices
+
+# Failed to get convolution algorithm
+physical_devices = list_physical_devices('GPU')
+set_memory_growth(physical_devices[0], True)
 
 # Hyperparameters
 from cycle.config import ORIG_IMG_SIZE
 from cycle.config import INPUT_IMG_SIZE
 from cycle.config import BUFFER_SIZE
 from cycle.config import BATCH_SIZE
 
+from cycle import CycleGAN
+from cycle.callbacks import GANMonitor
+from cycle.loss import generator_loss_fn
+from cycle.loss import discriminator_loss_fn
+from cycle.generator import get_resnet_generator
+from cycle.discriminator  import get_discriminator
+
 from cycle.preprocessing import preprocess_train_image
 from cycle.preprocessing import preprocess_test_image
 
@@ -64,4 +80,37 @@
     zebra = (((samples[1][0] * 127.5) + 127.5).numpy()).astype(np.uint8)
     ax_test[i, 0].imshow(horse)
     ax_test[i, 1].imshow(zebra)
-plt.show()
+plt.show()
+
+# Putting all together
+generator_G = get_resnet_generator(name='generator_G')
+generator_F = get_resnet_generator(name='generator_F')
+
+discriminator_X = get_discriminator(name='discriminator_X')
+discriminator_Y = get_discriminator(name='discriminator_Y')
+
+cycle_model = CycleGAN(
+                generator_G=generator_G,
+                generator_F=generator_F,
+                discriminator_X=discriminator_X,
+                discriminator_Y=discriminator_Y)
+
+cycle_model.compile(
+    generator_G_opt=Adam(learning_rate=2e-4, beta_1=0.5),
+    generator_F_opt=Adam(learning_rate=2e-4, beta_1=0.5),
+    discriminator_X_opt=Adam(learning_rate=2e-4, beta_1=0.5),
+    discriminator_Y_opt=Adam(learning_rate=2e-4, beta_1=0.5),
+    generator_loss_fn=generator_loss_fn,
+    discriminator_loss_fn=discriminator_loss_fn)
+
+plotter = GANMonitor(data=test_horses)
+checkpoint_filepath = "./model_checkpoints/cyclegan_checkpoints.{epoch:03d}"
+model_checkpoint_callback = ModelCheckpoint(
+    filepath=checkpoint_filepath
+)
+
+cycle_model.fit(
+    Dataset.zip((train_horses, train_zebras)),
+    epochs=90,
+    callbacks=[plotter, model_checkpoint_callback],
+)