chore: expose gs tau argument [SD-128] (#145)

* Expose tau argument of the Gumbel Softmax

* Integrate exposed tau arg in regular synths

* expose tau only on RegularModels

Author: jfsantos-ds

src/ydata_synthetic/synthesizers/gan.py

@@ -17,9 +17,9 @@
 from ydata_synthetic.synthesizers.saving_keras import make_keras_picklable
 
 _model_parameters = ['batch_size', 'lr', 'betas', 'layers_dim', 'noise_dim',
-                     'n_cols', 'seq_len', 'condition', 'n_critic', 'n_features']
+                     'n_cols', 'seq_len', 'condition', 'n_critic', 'n_features', 'tau_gs']
 _model_parameters_df = [128, 1e-4, (None, None), 128, 264,
-                        None, None, None, 1, None]
+                        None, None, None, 1, None, 0.2]
 
 _train_parameters = ['cache_prefix', 'label_dim', 'epochs', 'sample_interval', 'labels']
 
@@ -62,6 +62,8 @@ def __init__(
         self.data_dim = None
         self.layers_dim = model_parameters.layers_dim
         self.processor = None
+        if self.__MODEL__ in RegularModels.__members__:
+            self.tau = model_parameters.tau_gs
 
     # pylint: disable=E1101
     def __call__(self, inputs, **kwargs):

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

@@ -54,7 +54,7 @@ def label_col(self, data_label: Tuple[Union[DataFrame, array], str]):
     def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.generator = Generator(self.batch_size, self.num_classes). \
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info = activation_info)
+                        activation_info = activation_info, tau = self.tau)
 
         self.discriminator = Discriminator(self.batch_size, self.num_classes). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -200,7 +200,7 @@ def __init__(self, batch_size, num_classes):
         self.batch_size = batch_size
         self.num_classes = num_classes
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None, tau: Optional[float] = None):
         noise = Input(shape=input_shape, batch_size=self.batch_size)
         label = Input(shape=(1,), batch_size=self.batch_size, dtype='int32')
         label_embedding = Flatten()(Embedding(self.num_classes, 1)(label))
@@ -211,7 +211,7 @@ def build_model(self, input_shape, dim, data_dim, activation_info: Optional[Name
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info).call(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=[noise, label], outputs=x)
 
 

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

@@ -30,7 +30,7 @@ def __init__(self, model_parameters, gradient_penalty_weight=10):
     def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.generator = Generator(self.batch_size). \
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info=activation_info)
+                        activation_info=activation_info, tau = self.tau)
 
         self.critic = Critic(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -192,14 +192,14 @@ def __init__(self, batch_size):
         """Simple generator with dense feedforward layers."""
         self.batch_size = batch_size
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None, tau: Optional[float] = None):
         input_ = Input(shape=input_shape, batch_size=self.batch_size)
         x = Dense(dim, activation='relu')(input_)
         x = Dense(dim * 2, activation='relu')(x)
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info)(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=input_, outputs=x)
 
 class Critic(tf.keras.Model):

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

@@ -27,7 +27,7 @@ def define_gan(self, col_transform_info: Optional[NamedTuple] = None):
         # define generator/discriminator
         self.generator = Generator(self.batch_size). \
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info=col_transform_info)
+                        activation_info=col_transform_info, tau = self.tau)
 
         self.discriminator = Discriminator(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -177,12 +177,12 @@ class Generator(Model):
     def __init__(self, batch_size):
         self.batch_size = batch_size
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: NamedTuple = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: NamedTuple = None, tau: Optional[float] = None):
         input = Input(shape=input_shape, batch_size = self.batch_size)
         x = Dense(dim, kernel_initializer=initializers.TruncatedNormal(mean=0., stddev=0.5), activation='relu')(input)
         x = Dense(dim * 2, activation='relu')(x)
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info)(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=input, outputs=x)

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

@@ -23,7 +23,7 @@ def __init__(self, model_parameters):
     def define_gan(self, activation_info: Optional[NamedTuple]):
         self.generator = Generator(self.batch_size).\
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info = activation_info)
+                        activation_info = activation_info, tau = self.tau)
 
         self.discriminator = Discriminator(self.batch_size).\
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -131,14 +131,14 @@ class Generator(tf.keras.Model):
     def __init__(self, batch_size):
         self.batch_size=batch_size
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None, tau: Optional[float] = None):
         input= Input(shape=input_shape, batch_size=self.batch_size)
         x = Dense(dim, activation='relu')(input)
         x = Dense(dim * 2, activation='relu')(x)
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info)(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=input, outputs=x)
 
 class Discriminator(tf.keras.Model):

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

@@ -45,7 +45,7 @@ def wasserstein_loss(self, y_true, y_pred):
     def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.generator = Generator(self.batch_size). \
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info=activation_info)
+                        activation_info=activation_info, tau = self.tau)
 
         self.critic = Critic(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -155,14 +155,14 @@ class Generator(tf.keras.Model):
     def __init__(self, batch_size):
         self.batch_size = batch_size
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None, tau: Optional[float] = None):
         input = Input(shape=input_shape, batch_size=self.batch_size)
         x = Dense(dim, activation='relu')(input)
         x = Dense(dim * 2, activation='relu')(x)
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info)(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=input, outputs=x)
 
 class Critic(tf.keras.Model):

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

@@ -28,7 +28,7 @@ def __init__(self, model_parameters, n_critic, gradient_penalty_weight=10):
     def define_gan(self, activation_info: Optional[NamedTuple] = None):
         self.generator = Generator(self.batch_size). \
             build_model(input_shape=(self.noise_dim,), dim=self.layers_dim, data_dim=self.data_dim,
-                        activation_info=activation_info)
+                        activation_info=activation_info, tau = self.tau)
 
         self.critic = Critic(self.batch_size). \
             build_model(input_shape=(self.data_dim,), dim=self.layers_dim)
@@ -176,14 +176,14 @@ class Generator(tf.keras.Model):
     def __init__(self, batch_size):
         self.batch_size = batch_size
 
-    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None):
+    def build_model(self, input_shape, dim, data_dim, activation_info: Optional[NamedTuple] = None, tau: Optional[float] = None):
         input = Input(shape=input_shape, batch_size=self.batch_size)
         x = Dense(dim, activation='relu')(input)
         x = Dense(dim * 2, activation='relu')(x)
         x = Dense(dim * 4, activation='relu')(x)
         x = Dense(data_dim)(x)
         if activation_info:
-            x = GumbelSoftmaxActivation(activation_info)(x)
+            x = GumbelSoftmaxActivation(activation_info, tau=tau)(x)
         return Model(inputs=input, outputs=x)
 
 class Critic(tf.keras.Model):

src/ydata_synthetic/utils/gumbel_softmax.py

@@ -22,9 +22,14 @@ def gumbel_noise(shape: TensorShape) -> Tensor:
 
 @register_keras_serializable(package='Synthetic Data', name='GumbelSoftmaxLayer')
 class GumbelSoftmaxLayer(Layer):
-    "A Gumbel-Softmax layer implementation that should be stacked on top of a categorical feature logits."
+    """A Gumbel-Softmax layer implementation that should be stacked on top of a categorical feature logits.
 
-    def __init__(self, tau: float = 0.2, name: Optional[str] = None, **kwargs):
+    Arguments:
+            tau (float): Temperature parameter of the GS layer
+            name (Optional[str]): Name for a single categorical block
+    """
+
+    def __init__(self, tau: float, name: Optional[str] = None, **kwargs):
         super().__init__(name=name, **kwargs)
         self.tau = tau
 
@@ -54,11 +59,15 @@ class GumbelSoftmaxActivation(Layer):
         processor's pipelines in/out feature maps. For simplicity this object can be taken directly from the data \
         processor col_transform_info."""
 
-    def __init__(self, activation_info: NamedTuple, name: Optional[str] = None, **kwargs):
+    def __init__(self, activation_info: NamedTuple, name: Optional[str] = None, tau: Optional[float] = None, **kwargs):
         """Arguments:
             col_map (NamedTuple): Defines each of the processor pipelines input/output features.
-            name (Optional[str]): Name of the layer"""
+            name (Optional[str]): Name of the GumbelSoftmaxActivation layer
+            tau (Optional[float]): Temperature parameter of the GS layer, must be a float bigger than 0"""
         super().__init__(name=name, **kwargs)
+        self.tau = 0.2 if not tau else tau  # Defaults to the default value proposed in the original article
+        assert isinstance(self.tau, (int, float)) and self.tau > 0, "Optional argument tau must be numerical and \
+bigger than 0."
 
         self._activation_info = activation_info
 
@@ -74,7 +83,7 @@ def call(self, _input):  # pylint: disable=W0221
         cat_cols = split(cat_cols, self._cat_lens if self._cat_lens else [0], 1, name='split_cats')
 
         num_cols = [Activation('tanh', name='num_cols_activation')(num_cols)]
-        cat_cols = [GumbelSoftmaxLayer(name=name)(col)[0] for name, col in \
+        cat_cols = [GumbelSoftmaxLayer(tau=self.tau, name=name)(col)[0] for name, col in \
             zip(self.cat_feats.feat_names_in, cat_cols)]
         return concat(num_cols+cat_cols, 1)