fix: remove optimizers as model attributess (#160)

* remove optimizers from dragan

* remove  optimizers from cramergan

* remove optimizers from wgangp

* remove optimizers from cwgangp

Author: jfsantos-ds

src/ydata_synthetic/synthesizers/regular/cramergan/model.py

@@ -35,19 +35,21 @@ def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.critic = Critic(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
 
-        self.g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
-        self.c_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        c_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
 
         # The generator takes noise as input and generates records
         z = Input(shape=(self.noise_dim,), batch_size=self.batch_size)
         fake = self.generator(z)
         logits = self.critic(fake)
 
+        return g_optimizer, c_optimizer
+
     def gradient_penalty(self, real, fake):
         gp = gradient_penalty(self.f_crit, real, fake, mode=Mode.CRAMER)
         return gp
 
-    def update_gradients(self, x):
+    def update_gradients(self, x, g_optimizer, c_optimizer):
         """Compute and apply the gradients for both the Generator and the Critic.
 
         :param x: real data event
@@ -71,13 +73,13 @@ def update_gradients(self, x):
         g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
 
         # Update the weights of the generator
-        self.g_optimizer.apply_gradients(
+        g_optimizer.apply_gradients(
             zip(g_gradients, self.generator.trainable_variables)
         )
 
         c_gradient = d_tape.gradient(c_loss, self.critic.trainable_variables)
         # Update the weights of the critic using the optimizer
-        self.c_optimizer.apply_gradients(
+        c_optimizer.apply_gradients(
             zip(c_gradient, self.critic.trainable_variables)
         )
 
@@ -131,8 +133,8 @@ def get_data_batch(train, batch_size, seed=0):
         train_ix = list(train_ix) + list(train_ix)  # duplicate to cover ranges past the end of the set
         return train[train_ix[start_i: stop_i]]
 
-    def train_step(self, train_data):
-        critic_loss, g_loss = self.update_gradients(train_data)
+    def train_step(self, train_data, optimizers):
+        critic_loss, g_loss = self.update_gradients(train_data, *optimizers)
         return critic_loss, g_loss
 
     def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat_cols: List[str]):
@@ -147,7 +149,7 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
 
         data = self.processor.transform(data)
         self.data_dim = data.shape[1]
-        self.define_gan(self.processor.col_transform_info)
+        optimizers = self.define_gan(self.processor.col_transform_info)
 
         iterations = int(abs(data.shape[0] / self.batch_size) + 1)
 
@@ -158,7 +160,7 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
             for epoch in trange(train_arguments.epochs):
                 for iteration in range(iterations):
                     batch_data = self.get_data_batch(data, self.batch_size)
-                    c_loss, g_loss = self.train_step(batch_data)
+                    c_loss, g_loss = self.train_step(batch_data, optimizers)
 
                     if iteration % train_arguments.sample_interval == 0:
                         # Test here data generation step
@@ -168,23 +170,7 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
                         model_checkpoint_base_name = './cache/' + train_arguments.cache_prefix + '_{}_model_weights_step_{}.h5'
                         self.generator.save_weights(model_checkpoint_base_name.format('generator', iteration))
                         self.critic.save_weights(model_checkpoint_base_name.format('critic', iteration))
-
-                print(
-                    "Epoch: {} | critic_loss: {} | gen_loss: {}".format(
-                        epoch, c_loss, g_loss
-                    ))
-
-        self.g_optimizer=self.g_optimizer.get_config()
-        self.critic_optimizer=self.c_optimizer.get_config()
-
-    def save(self, path):
-        """Strip down the optimizers from the model then save."""
-        for attr in ['g_optimizer', 'c_optimizer']:
-            try:
-                delattr(self, attr)
-            except AttributeError:
-                continue
-        super().save(path)
+                print(f"Epoch: {epoch} | critic_loss: {c_loss} | gen_loss: {g_loss}")
 
 
 class Generator(tf.keras.Model):

src/ydata_synthetic/synthesizers/regular/cwgangp/model.py

@@ -45,8 +45,9 @@ def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.critic = Critic(self.batch_size, self.num_classes). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
 
-        self.g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
-        self.critic_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        c_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        return g_optimizer, c_optimizer
 
     def gradient_penalty(self, real, fake, label):
         epsilon = random.uniform([real.shape[0], 1], 0.0, 1.0, dtype=dtypes.float32)
@@ -130,7 +131,7 @@ def train(self, data: DataFrame, label_col: str, train_arguments: TrainParameter
 
         processed_data = self.processor.transform(data)
         self.data_dim = processed_data.shape[1]
-        self.define_gan(self.processor.col_transform_info)
+        optimizers = self.define_gan(self.processor.col_transform_info)
 
         # Merging labels with processed data
         processed_data = hstack([processed_data, label])
@@ -145,7 +146,7 @@ def train(self, data: DataFrame, label_col: str, train_arguments: TrainParameter
                 batch_x = self.get_data_batch(processed_data, self.batch_size)  # Batches are retrieved with labels
                 batch_x, label = batch_x[:, :-1], batch_x[:, -1]  # Separate labels from batch
 
-                cri_loss, ge_loss = self.train_step((batch_x, label))
+                cri_loss, ge_loss = self.train_step((batch_x, label), optimizers)
 
             print(
                 "Epoch: {} | critic_loss: {} | gen_loss: {}".format(

src/ydata_synthetic/synthesizers/regular/dragan/model.py

@@ -32,17 +32,19 @@ def define_gan(self, col_transform_info: Optional[NamedTuple] = None):
         self.discriminator = Discriminator(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
 
-        self.g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2, clipvalue=0.001)
-        self.d_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2, clipvalue=0.001)
+        g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2, clipvalue=0.001)
+        d_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2, clipvalue=0.001)
+        return g_optimizer, d_optimizer
 
     def gradient_penalty(self, real, fake):
         gp = gradient_penalty(self.discriminator, real, fake, mode= Mode.DRAGAN)
         return gp
 
-    def update_gradients(self, x):
+    def update_gradients(self, x, g_optimizer, d_optimizer):
         """
         Compute the gradients for both the Generator and the Discriminator
-        :param x: real data event
+            x (tf.tensor): real data event
+            *_optimizer (tf.OptimizerV2): Optimizer for the * model
         :return: generator gradients, discriminator gradients
         """
         # Update the gradients of critic for n_critic times (Training the critic)
@@ -52,7 +54,7 @@ def update_gradients(self, x):
             # Get the gradients of the critic
             d_gradient = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
             # Update the weights of the critic using the optimizer
-            self.d_optimizer.apply_gradients(
+            d_optimizer.apply_gradients(
                 zip(d_gradient, self.discriminator.trainable_variables)
             )
 
@@ -64,7 +66,7 @@ def update_gradients(self, x):
         gen_gradients = g_tape.gradient(gen_loss, self.generator.trainable_variables)
 
         # Update the weights of the generator
-        self.g_optimizer.apply_gradients(
+        g_optimizer.apply_gradients(
             zip(gen_gradients, self.generator.trainable_variables)
         )
 
@@ -112,8 +114,8 @@ def get_data_batch(self, train, batch_size):
             .batch(batch_size).shuffle(buffer_size)
         return train_loader
 
-    def train_step(self, train_data):
-        d_loss, g_loss = self.update_gradients(train_data)
+    def train_step(self, train_data, optimizers):
+        d_loss, g_loss = self.update_gradients(train_data, *optimizers)
         return d_loss, g_loss
 
     def train(self, data, train_arguments, num_cols, cat_cols):
@@ -128,7 +130,7 @@ def train(self, data, train_arguments, num_cols, cat_cols):
 
         processed_data = self.processor.transform(data)
         self.data_dim = processed_data.shape[1]
-        self.define_gan(self.processor.col_transform_info)
+        optimizers = self.define_gan(self.processor.col_transform_info)
 
         train_loader = self.get_data_batch(processed_data, self.batch_size)
 
@@ -139,7 +141,7 @@ def train(self, data, train_arguments, num_cols, cat_cols):
             for epoch in tqdm.trange(train_arguments.epochs):
                 for batch_data in train_loader:
                     batch_data = tf.cast(batch_data, dtype=tf.float32)
-                    d_loss, g_loss = self.train_step(batch_data)
+                    d_loss, g_loss = self.train_step(batch_data, optimizers)
 
                 print(
                     "Epoch: {} | disc_loss: {} | gen_loss: {}".format(
@@ -155,9 +157,6 @@ def train(self, data, train_arguments, num_cols, cat_cols):
                     self.generator.save_weights(model_checkpoint_base_name.format('generator', epoch))
                     self.discriminator.save_weights(model_checkpoint_base_name.format('discriminator', epoch))
 
-            self.g_optimizer=self.g_optimizer.get_config()
-            self.d_optimizer=self.d_optimizer.get_config()
-
 
 class Discriminator(Model):
     def __init__(self, batch_size):

src/ydata_synthetic/synthesizers/regular/wgangp/model.py

@@ -33,8 +33,9 @@ def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.critic = Critic(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
 
-        self.g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
-        self.critic_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        g_optimizer = Adam(self.g_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        c_optimizer = Adam(self.d_lr, beta_1=self.beta_1, beta_2=self.beta_2)
+        return g_optimizer, c_optimizer
 
     def gradient_penalty(self, real, fake):
         epsilon = tf.random.uniform([real.shape[0], 1], 0.0, 1.0, dtype=tf.dtypes.float32)
@@ -47,7 +48,7 @@ def gradient_penalty(self, real, fake):
         d_regularizer = tf.reduce_mean((ddx - 1.0) ** 2)
         return d_regularizer
 
-    def update_gradients(self, x):
+    def update_gradients(self, x, g_optimizer, c_optimizer):
         """
         Compute the gradients for both the Generator and the Critic
         :param x: real data event
@@ -60,7 +61,7 @@ def update_gradients(self, x):
             # Get the gradients of the critic
             d_gradient = d_tape.gradient(critic_loss, self.critic.trainable_variables)
             # Update the weights of the critic using the optimizer
-            self.critic_optimizer.apply_gradients(
+            c_optimizer.apply_gradients(
                 zip(d_gradient, self.critic.trainable_variables)
             )
 
@@ -72,7 +73,7 @@ def update_gradients(self, x):
         gen_gradients = g_tape.gradient(gen_loss, self.generator.trainable_variables)
 
         # Update the weights of the generator
-        self.g_optimizer.apply_gradients(
+        g_optimizer.apply_gradients(
             zip(gen_gradients, self.generator.trainable_variables)
         )
 
@@ -124,8 +125,8 @@ def get_data_batch(self, train, batch_size, seed=0):
         return train[train_ix[start_i: stop_i]]
 
     @tf.function
-    def train_step(self, train_data):
-        cri_loss, ge_loss = self.update_gradients(train_data)
+    def train_step(self, train_data, optimizers):
+        cri_loss, ge_loss = self.update_gradients(train_data, *optimizers)
         return cri_loss, ge_loss
 
     def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat_cols: List[str]):
@@ -140,7 +141,7 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
 
         processed_data = self.processor.transform(data)
         self.data_dim = processed_data.shape[1]
-        self.define_gan(self.processor.col_transform_info)
+        optimizers = self.define_gan(self.processor.col_transform_info)
 
         iterations = int(abs(data.shape[0]/self.batch_size)+1)
 
@@ -151,7 +152,7 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
             for epoch in trange(train_arguments.epochs):
                 for _ in range(iterations):
                     batch_data = self.get_data_batch(processed_data, self.batch_size).astype(np.float32)
-                    cri_loss, ge_loss = self.train_step(batch_data)
+                    cri_loss, ge_loss = self.train_step(batch_data, optimizers)
 
                 print(
                     "Epoch: {} | disc_loss: {} | gen_loss: {}".format(
@@ -167,9 +168,6 @@ def train(self, data, train_arguments: TrainParameters, num_cols: List[str], cat
                     self.generator.save_weights(model_checkpoint_base_name.format('generator', epoch))
                     self.critic.save_weights(model_checkpoint_base_name.format('critic', epoch))
 
-        self.g_optimizer=self.g_optimizer.get_config()
-        self.critic_optimizer=self.critic_optimizer.get_config()
-
 
 class Generator(tf.keras.Model):
     def __init__(self, batch_size):