reduc tf logging when doing hyper search

ppdebreuck · Nov 13, 2021 · 6a5b139 · 6a5b139
1 parent 82e83f9
commit 6a5b139
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/modnet/hyper_opt/fit_genetic.py b/modnet/hyper_opt/fit_genetic.py
@@ -308,6 +308,11 @@ def function_fitness(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         num_nested_folds = 5
         if nested:
@@ -381,6 +386,8 @@ def function_fitness(
         pool.close()
         pool.join()
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return val_loss_per_individual, np.array(models), np.array(individuals)
 
     def run(

diff --git a/modnet/models/ensemble.py b/modnet/models/ensemble.py
@@ -244,6 +244,11 @@ def fit_preset(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         if callbacks is None:
             es = tf.keras.callbacks.EarlyStopping(
@@ -403,6 +408,8 @@ def fit_preset(
                     final_models += models[idx][i].model
             self.__init__(modnet_models=final_models)
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return models, val_losses, best_learning_curve, learning_curves, best_preset
 
     def _make_picklable(self):

diff --git a/modnet/models/vanilla.py b/modnet/models/vanilla.py
@@ -411,6 +411,11 @@ def fit_preset(
         """
 
         from modnet.matbench.benchmark import matbench_kfold_splits
+        import os
+
+        os.environ[
+            "TF_CPP_MIN_LOG_LEVEL"
+        ] = "2"  # many models will be fitted => reduce output
 
         if callbacks is None:
             es = tf.keras.callbacks.EarlyStopping(
@@ -559,6 +564,8 @@ def fit_preset(
             self.model = best_model.model
             self._scaler = best_model._scaler
 
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"  # reset
+
         return models, val_losses, best_learning_curve, learning_curves, best_preset
 
     def predict(self, test_data: MODData, return_prob=False) -> pd.DataFrame: