load trained models

simplegan/gan/cgan.py

@@ -55,6 +55,8 @@ def __init__(
         activation="relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.image_size = None
@@ -278,6 +280,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,
@@ -432,6 +441,9 @@ def generate_samples(self, n_samples=1, labels_list=None, save_dir=None):
             len(labels_list) == n_samples
         ), "Number of samples does not match length of labels list"
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = np.random.uniform(-1, 1, (n_samples, self.noise_dim))
         labels_list = np.array(labels_list)
         generated_samples = self.gen_model([Z, labels_list]).numpy()

simplegan/gan/cyclegan.py

@@ -58,6 +58,10 @@ def __init__(
         disc_channels=[64, 128, 256, 512],
         kernel_size=(4, 4),
         kernel_initializer=tf.random_normal_initializer(0.0, 0.02),
+        gen_g_path=None,
+        gen_f_path=None,
+        disc_x_path=None,
+        disc_y_path=None,
     ):
 
         Pix2Pix.__init__(
@@ -75,6 +79,8 @@ def __init__(
         self.disc_model_x = None
         self.disc_model_y = None
 
+        self.config = locals()
+
     def load_data(
         self,
         data_dir=None,
@@ -280,6 +286,19 @@ def __load_model(self):
             self.discriminator(),
         )
 
+        if self.config["gen_g_path"] is not None:
+            self.gen_model_g.load_weights(self.config["gen_g_path"])
+            print("Generator-G checkpoint restored")
+        if self.config["gen_f_path"] is not None:
+            self.gen_model_f.load_weights(self.config["gen_f_path"])
+            print("Generator-F checkpoint restored")
+        if self.config["disc_x_path"] is not None:
+            self.disc_model_x.load_weights(self.config["disc_x_path"])
+            print("Discriminator-X checkpoint restored")
+        if self.config["disc_y_path"] is not None:
+            self.disc_model_y.load_weights(self.config["disc_y_path"])
+            print("Discriminator-Y checkpoint restored")
+
     def _save_samples(self, model, image, count):
 
         assert os.path.exists(self.save_img_dir), "sample directory does not exist"
@@ -552,8 +571,18 @@ def fit(
             assert isinstance(save_model, str), "Not a valid directory"
             if save_model[-1] != "/":
                 self.gen_model_g.save_weights(save_model + "/generator_g_checkpoint")
+                self.gen_model_f.save_weights(save_model + "/generator_f_checkpoint")
+                self.disc_model_x.save_weights(
+                    save_model + "/discrimnator_x_checkpoint"
+                )
+                self.disc_model_y.save_weights(
+                    save_model + "/discrimnator_y_checkpoint"
+                )
             else:
                 self.gen_model_g.save_weights(save_model + "generator_g_checkpoint")
+                self.gen_model_f.save_weights(save_model + "generator_f_checkpoint")
+                self.disc_model_x.save_weights(save_model + "discrimnator_x_checkpoint")
+                self.disc_model_y.save_weights(save_model + "discrimnator_y_checkpoint")
 
     def generate_samples(self, test_ds=None, save_dir=None):
 
@@ -569,6 +598,9 @@ def generate_samples(self, test_ds=None, save_dir=None):
 
         assert test_ds is not None, "Enter input test dataset"
 
+        if self.gen_model_g is None:
+            self.__load_model()
+
         generated_samples = []
         for image in test_ds:
             gen_image = self.gen_model_g(image, training=False).numpy()

simplegan/gan/dcgan.py

@@ -57,6 +57,8 @@ def __init__(
         activation="relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.image_size = None
@@ -109,7 +111,7 @@ def load_data(
 
         else:
 
-            train_data = load_custom_data(data_dir, img_size)
+            train_data = load_custom_data(data_dir, img_shape)
 
         self.image_size = train_data.shape[1:]
 
@@ -321,6 +323,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,
@@ -466,6 +475,9 @@ def generate_samples(self, n_samples=1, save_dir=None):
             returns ``None`` if save_dir is ``not None``, otherwise returns a numpy array with generated samples
         """
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = tf.random.normal([n_samples, self.noise_dim])
         generated_samples = self.gen_model(Z).numpy()
 

simplegan/gan/infogan.py

@@ -5,7 +5,7 @@
 from tensorflow.keras import layers
 from ..datasets.load_mnist import load_mnist
 from ..datasets.load_cifar10 import load_cifar10
-from ..datasets.load_custom_data import load_custom_data
+from ..datasets.load_custom_data import load_custom_data_with_labels
 from ..losses.minmax_loss import gan_discriminator_loss, gan_generator_loss
 from ..losses.infogan_loss import auxillary_loss
 import datetime
@@ -57,6 +57,8 @@ def __init__(
         activation="leaky_relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.image_size = None
@@ -304,6 +306,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,
@@ -457,6 +466,9 @@ def generate_samples(self, n_samples=1, save_dir=None):
             returns ``None`` if save_dir is ``not None``, otherwise returns a numpy array with generated samples
         """
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = np.random.randn(n_samples, self.noise_dim)
         label_input = tf.keras.utils.to_categorical(
             (np.random.randint(0, self.n_classes, n_samples)), self.n_classes

simplegan/gan/pix2pix.py

@@ -56,6 +56,8 @@ def __init__(
         disc_channels=[64, 128, 256, 512],
         kernel_size=(4, 4),
         kernel_initializer=tf.random_normal_initializer(0.0, 0.02),
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.gen_model = None
@@ -365,6 +367,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def _save_samples(self, model, ex_input, ex_target, count):
 
         assert os.path.exists(self.save_img_dir), "sample directory does not exist"
@@ -582,6 +591,9 @@ def generate_samples(self, test_ds=None, save_dir=None):
 
         assert test_ds is not None, "Enter input test dataset"
 
+        if self.gen_model is None:
+            self.__load_model()
+
         generated_samples = []
         for image in test_ds:
             gen_image = self.gen_model(image, training=False).numpy()

simplegan/gan/sagan.py

@@ -115,13 +115,13 @@ class SAGAN:
         noise_dim (int, optional): represents the dimension of the prior to sample values. Defaults to ``128``
     """
 
-    def __init__(
-        self, noise_dim=128,
-    ):
+    def __init__(self, noise_dim=128, gen_path=None, disc_path=None):
 
         self.image_size = None
         self.noise_dim = noise_dim
         self.n_classes = None
+        self.gen_model = None
+        self.disc_model = None
         self.config = locals()
 
     def load_data(
@@ -211,6 +211,13 @@ def __load_model(self):
             Discriminator(self.n_classes),
         )
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     @tf.function
     def train_step(self, images, labels):
 
@@ -381,6 +388,9 @@ def generate_samples(self, n_samples=1, labels_list=None, save_dir=None):
             len(labels_list) == n_samples
         ), "Number of samples does not match length of labels list"
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = np.random.uniform(-1, 1, (n_samples, self.noise_dim))
         labels_list = np.array(labels_list)
         generated_samples = self.gen_model([Z, labels_list]).numpy()

simplegan/gan/vanilla_gan.py

@@ -52,6 +52,8 @@ def __init__(
         activation="relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.image_size = None
@@ -241,6 +243,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,
@@ -380,6 +389,9 @@ def generate_samples(self, n_samples=1, save_dir=None):
             returns ``None`` if save_dir is ``not None``, otherwise returns a numpy array with generated samples
         """
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = np.random.uniform(-1, 1, (n_samples, self.noise_dim))
         generated_samples = self.gen_model(Z)
         generated_samples = tf.reshape(

simplegan/gan/voxelgan.py

@@ -44,6 +44,8 @@ def __init__(
         activation="relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         self.noise_dim = noise_dim
@@ -256,6 +258,13 @@ def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,
@@ -407,6 +416,9 @@ def generate_sample(self, n_samples=1, plot=False):
             ``None`` if ``plot`` is ``True`` else a numpy array of samples of shape ``(n_samples, side_length, side_length, side_length, 1)``
         """
 
+        if self.gen_model is None:
+            self.__load_model()
+
         Z = np.random.uniform(0, 1, (n_samples, self.noise_dim))
         generated_samples = self.gen_model(Z).numpy()
 

simplegan/gan/wgan.py

@@ -58,6 +58,8 @@ def __init__(
         activation="relu",
         kernel_initializer="glorot_uniform",
         kernel_regularizer=None,
+        gen_path=None,
+        disc_path=None,
     ):
 
         DCGAN.__init__(
@@ -70,12 +72,21 @@ def __init__(
             activation,
             kernel_initializer,
             kernel_regularizer,
+            gen_path,
+            disc_path,
         )
 
     def __load_model(self):
 
         self.gen_model, self.disc_model = self.generator(), self.discriminator()
 
+        if self.config["gen_path"] is not None:
+            self.gen_model.load_weights(self.config["gen_path"])
+            print("Generator checkpoint restored")
+        if self.config["disc_path"] is not None:
+            self.disc_model.load_weights(self.config["disc_path"])
+            print("Discriminator checkpoint restored")
+
     def fit(
         self,
         train_ds=None,