updates for issue #10

mbarbetti · Jul 19, 2024 · 286c313 · 286c313
1 parent 8eda454
commit 286c313
Show file tree

Hide file tree

Showing 17 changed files with 171 additions and 83 deletions.
diff --git a/scripts/train_ANN_isMuon.py b/scripts/train_ANN_isMuon.py
@@ -1,4 +1,6 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import yaml
 import pickle
 import keras as k
@@ -20,8 +22,6 @@
 from pidgan.utils.preprocessing import invertColumnTransformer
 from pidgan.utils.reports import initHPSingleton
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 DTYPE = np.float32
 BATCHSIZE = 2048
 

diff --git a/scripts/train_GAN_GlobalPID-im.py b/scripts/train_GAN_GlobalPID-im.py
@@ -1,4 +1,6 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import yaml
 import pickle
 import keras as k
@@ -23,8 +25,6 @@
 from pidgan.utils.preprocessing import invertColumnTransformer
 from pidgan.utils.reports import initHPSingleton
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 DTYPE = np.float32
 BATCHSIZE = 2048
 

diff --git a/scripts/train_GAN_GlobalPID-nm.py b/scripts/train_GAN_GlobalPID-nm.py
@@ -1,4 +1,6 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import yaml
 import pickle
 import keras as k
@@ -23,8 +25,6 @@
 from pidgan.utils.preprocessing import invertColumnTransformer
 from pidgan.utils.reports import initHPSingleton
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 DTYPE = np.float32
 BATCHSIZE = 2048
 

diff --git a/scripts/train_GAN_Muon.py b/scripts/train_GAN_Muon.py
@@ -1,4 +1,6 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import yaml
 import pickle
 import keras as k
@@ -23,8 +25,6 @@
 from pidgan.utils.preprocessing import invertColumnTransformer
 from pidgan.utils.reports import initHPSingleton
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 DTYPE = np.float32
 BATCHSIZE = 2048
 

diff --git a/scripts/train_GAN_Rich.py b/scripts/train_GAN_Rich.py
@@ -1,4 +1,6 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import yaml
 import pickle
 import keras as k
@@ -23,8 +25,6 @@
 from pidgan.utils.preprocessing import invertColumnTransformer
 from pidgan.utils.reports import initHPSingleton
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 DTYPE = np.float32
 BATCHSIZE = 2048
 

diff --git a/src/pidgan/algorithms/k3/BceGAN_ALP.py b/src/pidgan/algorithms/k3/BceGAN_ALP.py
@@ -41,6 +41,7 @@ def compile(
         virtual_adv_direction_upds=1,
         referee_optimizer=None,
         referee_upds_per_batch=None,
+        jit_compile=False,
     ) -> None:
         super().compile(
             metrics=metrics,
@@ -50,6 +51,7 @@ def compile(
             discriminator_upds_per_batch=discriminator_upds_per_batch,
             referee_optimizer=referee_optimizer,
             referee_upds_per_batch=referee_upds_per_batch,
+            jit_compile=jit_compile,
         )
 
         # Virtual adversarial direction updates

diff --git a/src/pidgan/algorithms/k3/GAN.py b/src/pidgan/algorithms/k3/GAN.py
@@ -96,8 +96,9 @@ def compile(
         discriminator_upds_per_batch=1,
         referee_optimizer=None,
         referee_upds_per_batch=None,
+        jit_compile=False,
     ) -> None:
-        super().compile(weighted_metrics=[])
+        super().compile(weighted_metrics=[], jit_compile=jit_compile)
 
         # Metrics
         if not self._model_is_built:

diff --git a/src/pidgan/algorithms/k3/WGAN_ALP.py b/src/pidgan/algorithms/k3/WGAN_ALP.py
@@ -41,6 +41,7 @@ def compile(
         virtual_adv_direction_upds=1,
         referee_optimizer=None,
         referee_upds_per_batch=None,
+        jit_compile=False,
     ) -> None:
         super().compile(
             metrics=metrics,
@@ -50,6 +51,7 @@ def compile(
             discriminator_upds_per_batch=discriminator_upds_per_batch,
             referee_optimizer=referee_optimizer,
             referee_upds_per_batch=referee_upds_per_batch,
+            jit_compile=jit_compile,
         )
 
         # Virtual adversarial direction updates

diff --git a/tests/algorithms/test_BceGAN.py b/tests/algorithms/test_BceGAN.py
@@ -1,22 +1,22 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import pytest
 import warnings
-import keras as k
+import numpy as np
 import tensorflow as tf
 
 from pidgan.players.classifiers import AuxClassifier
 from pidgan.players.discriminators import AuxDiscriminator
 from pidgan.players.generators import ResGenerator
 from pidgan.metrics import BinaryCrossentropy as BCE
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 CHUNK_SIZE = int(1e4)
 BATCH_SIZE = 500
 
-x = tf.random.normal(shape=(CHUNK_SIZE, 4))
-y = tf.random.normal(shape=(CHUNK_SIZE, 8))
-w = tf.random.uniform(shape=(CHUNK_SIZE,))
+x = np.random.normal(size=(CHUNK_SIZE, 4)).astype("float32")
+y = np.random.normal(size=(CHUNK_SIZE, 8)).astype("float32")
+w = np.random.uniform(size=(CHUNK_SIZE,)).astype("float32")
 
 gen = ResGenerator(
     output_dim=y.shape[1],
@@ -120,6 +120,8 @@ def test_model_use(referee):
 @pytest.mark.parametrize("build_first", [True, False])
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 def test_model_compilation(model, build_first, metrics):
+    import keras as k
+
     if build_first:
         model(x, y)  # to build the model
 
@@ -169,6 +171,7 @@ def test_model_compilation(model, build_first, metrics):
 @pytest.mark.parametrize("sample_weight", [w, None])
 @pytest.mark.parametrize("build_first", [True, False])
 def test_model_train(referee, sample_weight, build_first):
+    import keras as k
     from pidgan.algorithms import BceGAN
 
     if sample_weight is not None:
@@ -236,14 +239,21 @@ def test_model_train(referee, sample_weight, build_first):
 
     for s in states:
         for entry in train.history[s]:
-            print(train.history)
-            print(f"{s}: {entry}")
             assert isinstance(entry, (int, float))
 
+    if not build_first:
+        for key in ["g_loss", "d_loss", "bce"]:
+            comparison = (
+                np.array(train.history[key]) != np.array(train.history[f"val_{key}"])
+            )
+            assert comparison.all()
+
 
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 @pytest.mark.parametrize("sample_weight", [w, None])
 def test_model_eval(model, metrics, sample_weight):
+    import keras as k
+
     g_opt = k.optimizers.RMSprop(learning_rate=0.001)
     d_opt = k.optimizers.RMSprop(learning_rate=0.001)
     r_opt = k.optimizers.RMSprop(learning_rate=0.001)

diff --git a/tests/algorithms/test_BceGAN_ALP.py b/tests/algorithms/test_BceGAN_ALP.py
@@ -1,22 +1,22 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import pytest
 import warnings
-import keras as k
+import numpy as np
 import tensorflow as tf
 
 from pidgan.players.classifiers import AuxClassifier
 from pidgan.players.discriminators import AuxDiscriminator
 from pidgan.players.generators import ResGenerator
 from pidgan.metrics import BinaryCrossentropy as BCE
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 CHUNK_SIZE = int(1e4)
 BATCH_SIZE = 500
 
-x = tf.random.normal(shape=(CHUNK_SIZE, 4))
-y = tf.random.normal(shape=(CHUNK_SIZE, 8))
-w = tf.random.uniform(shape=(CHUNK_SIZE,))
+x = np.random.normal(size=(CHUNK_SIZE, 4)).astype("float32")
+y = np.random.normal(size=(CHUNK_SIZE, 8)).astype("float32")
+w = np.random.uniform(size=(CHUNK_SIZE,)).astype("float32")
 
 gen = ResGenerator(
     output_dim=y.shape[1],
@@ -117,6 +117,8 @@ def test_model_use(referee):
 @pytest.mark.parametrize("build_first", [True, False])
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 def test_model_compilation(model, build_first, metrics):
+    import keras as k
+
     if build_first:
         model(x, y)  # to build the model
 
@@ -169,6 +171,7 @@ def test_model_compilation(model, build_first, metrics):
 @pytest.mark.parametrize("lipschitz_penalty_strategy", ["two-sided", "one-sided"])
 @pytest.mark.parametrize("build_first", [True, False])
 def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_first):
+    import keras as k
     from pidgan.algorithms import BceGAN_ALP
 
     if sample_weight is not None:
@@ -236,14 +239,21 @@ def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_f
 
     for s in states:
         for entry in train.history[s]:
-            print(train.history)
-            print(f"{s}: {entry}")
             assert isinstance(entry, (int, float))
 
+    if not build_first:
+        for key in ["g_loss", "d_loss", "bce"]:
+            comparison = (
+                np.array(train.history[key]) != np.array(train.history[f"val_{key}"])
+            )
+            assert comparison.all()
+
 
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 @pytest.mark.parametrize("sample_weight", [w, None])
 def test_model_eval(model, metrics, sample_weight):
+    import keras as k
+
     g_opt = k.optimizers.RMSprop(learning_rate=0.001)
     d_opt = k.optimizers.RMSprop(learning_rate=0.001)
     r_opt = k.optimizers.RMSprop(learning_rate=0.001)

diff --git a/tests/algorithms/test_BceGAN_GP.py b/tests/algorithms/test_BceGAN_GP.py
@@ -1,22 +1,22 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import pytest
 import warnings
-import keras as k
+import numpy as np
 import tensorflow as tf
 
 from pidgan.players.classifiers import AuxClassifier
 from pidgan.players.discriminators import AuxDiscriminator
 from pidgan.players.generators import ResGenerator
 from pidgan.metrics import BinaryCrossentropy as BCE
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 CHUNK_SIZE = int(1e4)
 BATCH_SIZE = 500
 
-x = tf.random.normal(shape=(CHUNK_SIZE, 4))
-y = tf.random.normal(shape=(CHUNK_SIZE, 8))
-w = tf.random.uniform(shape=(CHUNK_SIZE,))
+x = np.random.normal(size=(CHUNK_SIZE, 4)).astype("float32")
+y = np.random.normal(size=(CHUNK_SIZE, 8)).astype("float32")
+w = np.random.uniform(size=(CHUNK_SIZE,)).astype("float32")
 
 gen = ResGenerator(
     output_dim=y.shape[1],
@@ -117,6 +117,8 @@ def test_model_use(referee):
 @pytest.mark.parametrize("build_first", [True, False])
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 def test_model_compilation(model, build_first, metrics):
+    import keras as k
+
     if build_first:
         model(x, y)  # to build the model
 
@@ -167,6 +169,7 @@ def test_model_compilation(model, build_first, metrics):
 @pytest.mark.parametrize("lipschitz_penalty_strategy", ["two-sided", "one-sided"])
 @pytest.mark.parametrize("build_first", [True, False])
 def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_first):
+    import keras as k
     from pidgan.algorithms import BceGAN_GP
 
     if sample_weight is not None:
@@ -233,14 +236,21 @@ def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_f
 
     for s in states:
         for entry in train.history[s]:
-            print(train.history)
-            print(f"{s}: {entry}")
             assert isinstance(entry, (int, float))
 
+    if not build_first:
+        for key in ["g_loss", "d_loss", "bce"]:
+            comparison = (
+                np.array(train.history[key]) != np.array(train.history[f"val_{key}"])
+            )
+            assert comparison.all()
+
 
 @pytest.mark.parametrize("metrics", [["bce"], [BCE()], None])
 @pytest.mark.parametrize("sample_weight", [w, None])
 def test_model_eval(model, metrics, sample_weight):
+    import keras as k
+
     g_opt = k.optimizers.RMSprop(learning_rate=0.001)
     d_opt = k.optimizers.RMSprop(learning_rate=0.001)
     r_opt = k.optimizers.RMSprop(learning_rate=0.001)

diff --git a/tests/algorithms/test_CramerGAN.py b/tests/algorithms/test_CramerGAN.py
@@ -1,22 +1,22 @@
 import os
+os.environ["KERAS_BACKEND"] = "tensorflow"
+
 import pytest
 import warnings
-import keras as k
+import numpy as np
 import tensorflow as tf
 
 from pidgan.players.classifiers import AuxClassifier
 from pidgan.players.discriminators import AuxDiscriminator
 from pidgan.players.generators import ResGenerator
 from pidgan.metrics import WassersteinDistance as Wass_dist
 
-os.environ["KERAS_BACKEND"] = "tensorflow"
-
 CHUNK_SIZE = int(1e4)
 BATCH_SIZE = 500
 
-x = tf.random.normal(shape=(CHUNK_SIZE, 4))
-y = tf.random.normal(shape=(CHUNK_SIZE, 8))
-w = tf.random.uniform(shape=(CHUNK_SIZE,))
+x = np.random.normal(size=(CHUNK_SIZE, 4)).astype("float32")
+y = np.random.normal(size=(CHUNK_SIZE, 8)).astype("float32")
+w = np.random.uniform(size=(CHUNK_SIZE,)).astype("float32")
 
 gen = ResGenerator(
     output_dim=y.shape[1],
@@ -118,6 +118,8 @@ def test_model_use(referee):
 @pytest.mark.parametrize("build_first", [True, False])
 @pytest.mark.parametrize("metrics", [["bce"], [Wass_dist()], None])
 def test_model_compilation(model, build_first, metrics):
+    import keras as k
+
     if build_first:
         model(x, y)  # to build the model
 
@@ -168,6 +170,7 @@ def test_model_compilation(model, build_first, metrics):
 @pytest.mark.parametrize("lipschitz_penalty_strategy", ["two-sided", "one-sided"])
 @pytest.mark.parametrize("build_first", [True, False])
 def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_first):
+    import keras as k
     from pidgan.algorithms import CramerGAN
 
     if sample_weight is not None:
@@ -233,14 +236,14 @@ def test_model_train(referee, sample_weight, lipschitz_penalty_strategy, build_f
 
     for s in states:
         for entry in train.history[s]:
-            print(train.history)
-            print(f"{s}: {entry}")
             assert isinstance(entry, (int, float))
 
 
 @pytest.mark.parametrize("metrics", [["wass_dist"], [Wass_dist()], None])
 @pytest.mark.parametrize("sample_weight", [w, None])
 def test_model_eval(model, metrics, sample_weight):
+    import keras as k
+
     g_opt = k.optimizers.RMSprop(learning_rate=0.001)
     d_opt = k.optimizers.RMSprop(learning_rate=0.001)
     r_opt = k.optimizers.RMSprop(learning_rate=0.001)